mirror of
https://github.com/Dvorinka/swingmusic-extended.git
synced 2026-06-05 04:53:01 +00:00
support in_quotes search query
This commit is contained in:
+158
-46
@@ -1,21 +1,23 @@
|
||||
"""
|
||||
This library contains all the functions related to the search functionality.
|
||||
"""
|
||||
from typing import List, Generator, TypeVar, Any
|
||||
import itertools
|
||||
from typing import Any, Generator, List, TypeVar
|
||||
|
||||
from rapidfuzz import fuzz, process
|
||||
from rapidfuzz import process, utils
|
||||
from unidecode import unidecode
|
||||
|
||||
from app import models
|
||||
from app.utils.remove_duplicates import remove_duplicates
|
||||
|
||||
from app.models.enums import FavType
|
||||
from app.models.track import Track
|
||||
from app.store.albums import AlbumStore
|
||||
from app.store.artists import ArtistStore
|
||||
from app.store.tracks import TrackStore
|
||||
|
||||
ratio = fuzz.ratio
|
||||
wratio = fuzz.WRatio
|
||||
from app.utils.remove_duplicates import remove_duplicates
|
||||
from app.db.sqlite.favorite import SQLiteFavoriteMethods as favdb
|
||||
|
||||
# ratio = fuzz.ratio
|
||||
# wratio = fuzz.WRatio
|
||||
|
||||
|
||||
class Cutoff:
|
||||
@@ -56,6 +58,7 @@ class SearchTracks:
|
||||
track_titles,
|
||||
score_cutoff=Cutoff.tracks,
|
||||
limit=Limit.tracks,
|
||||
processor=utils.default_process,
|
||||
)
|
||||
|
||||
tracks = [self.tracks[i[2]] for i in results]
|
||||
@@ -67,7 +70,7 @@ class SearchArtists:
|
||||
self.query = query
|
||||
self.artists = ArtistStore.artists
|
||||
|
||||
def __call__(self) -> list:
|
||||
def __call__(self):
|
||||
"""
|
||||
Gets all artists with a given name.
|
||||
"""
|
||||
@@ -78,6 +81,7 @@ class SearchArtists:
|
||||
artists,
|
||||
score_cutoff=Cutoff.artists,
|
||||
limit=Limit.artists,
|
||||
processor=utils.default_process,
|
||||
)
|
||||
|
||||
return [self.artists[i[2]] for i in results]
|
||||
@@ -100,17 +104,11 @@ class SearchAlbums:
|
||||
albums,
|
||||
score_cutoff=Cutoff.albums,
|
||||
limit=Limit.albums,
|
||||
processor=utils.default_process,
|
||||
)
|
||||
|
||||
return [self.albums[i[2]] for i in results]
|
||||
|
||||
# get all artists that matched the query
|
||||
# for get all albums from the artists
|
||||
# get all albums that matched the query
|
||||
# return [**artist_albums **albums]
|
||||
|
||||
# recheck next and previous artist on play next or add to playlist
|
||||
|
||||
|
||||
class SearchPlaylists:
|
||||
def __init__(self, playlists: List[models.Playlist], query: str) -> None:
|
||||
@@ -124,12 +122,13 @@ class SearchPlaylists:
|
||||
playlists,
|
||||
score_cutoff=Cutoff.playlists,
|
||||
limit=Limit.playlists,
|
||||
processor=utils.default_process,
|
||||
)
|
||||
|
||||
return [self.playlists[i[2]] for i in results]
|
||||
|
||||
|
||||
_type = List[models.Track | models.Album | models.Artist]
|
||||
_type = models.Track | models.Album | models.Artist
|
||||
_S2 = TypeVar("_S2")
|
||||
_ResultType = int | float
|
||||
|
||||
@@ -148,7 +147,7 @@ def get_titles(items: _type):
|
||||
yield text
|
||||
|
||||
|
||||
class SearchAll:
|
||||
class TopResults:
|
||||
"""
|
||||
Joins all tracks, albums and artists
|
||||
then fuzzy searches them as a single unit.
|
||||
@@ -156,11 +155,11 @@ class SearchAll:
|
||||
|
||||
@staticmethod
|
||||
def collect_all():
|
||||
all_items: _type = []
|
||||
all_items: list[_type] = []
|
||||
|
||||
all_items.extend(ArtistStore.artists)
|
||||
all_items.extend(TrackStore.tracks)
|
||||
all_items.extend(AlbumStore.albums)
|
||||
all_items.extend(ArtistStore.artists)
|
||||
|
||||
return all_items, get_titles(all_items)
|
||||
|
||||
@@ -169,44 +168,157 @@ class SearchAll:
|
||||
items = list(items)
|
||||
|
||||
results = process.extract(
|
||||
query=query,
|
||||
choices=items,
|
||||
score_cutoff=Cutoff.tracks,
|
||||
limit=20
|
||||
query=query, choices=items, score_cutoff=Cutoff.tracks, limit=1
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def sort_results(items: _type):
|
||||
def map_with_type(item: _type):
|
||||
"""
|
||||
Separates results into differrent lists using itertools.groupby.
|
||||
Map the results to their respective types.
|
||||
"""
|
||||
mapped_items = [
|
||||
{"type": "track", "item": item} if isinstance(item, models.Track) else
|
||||
{"type": "album", "item": item} if isinstance(item, models.Album) else
|
||||
{"type": "artist", "item": item} if isinstance(item, models.Artist) else
|
||||
{"type": "Unknown", "item": item} for item in items
|
||||
]
|
||||
if isinstance(item, models.Track):
|
||||
return {"type": "track", "item": item}
|
||||
|
||||
mapped_items.sort(key=lambda x: x["type"])
|
||||
if isinstance(item, models.Album):
|
||||
tracks = TrackStore.get_tracks_by_albumhash(item.albumhash)
|
||||
tracks = remove_duplicates(tracks)
|
||||
|
||||
groups = [
|
||||
list(group) for key, group in
|
||||
itertools.groupby(mapped_items, lambda x: x["type"])
|
||||
]
|
||||
item.get_date_from_tracks(tracks)
|
||||
try:
|
||||
item.duration = sum((t.duration for t in tracks))
|
||||
except AttributeError:
|
||||
item.duration = 0
|
||||
|
||||
# merge items of a group into a dict that looks like: {"albums": [album1, ...]}
|
||||
groups = [
|
||||
{f"{group[0]['type']}s": [i['item'] for i in group]} for group in groups
|
||||
]
|
||||
item.check_is_single(tracks)
|
||||
|
||||
return groups
|
||||
if not item.is_single:
|
||||
item.check_type()
|
||||
|
||||
item.is_favorite = favdb.check_is_favorite(
|
||||
item.albumhash, fav_type=FavType.album
|
||||
)
|
||||
|
||||
return {"type": "album", "item": item}
|
||||
|
||||
if isinstance(item, models.Artist):
|
||||
track_count = 0
|
||||
duration = 0
|
||||
|
||||
for track in TrackStore.get_tracks_by_artisthash(item.artisthash):
|
||||
track_count += 1
|
||||
duration += track.duration
|
||||
|
||||
album_count = AlbumStore.count_albums_by_artisthash(item.artisthash)
|
||||
|
||||
item.set_trackcount(track_count)
|
||||
item.set_albumcount(album_count)
|
||||
item.set_duration(duration)
|
||||
|
||||
return {"type": "artist", "item": item}
|
||||
|
||||
@staticmethod
|
||||
def search(query: str):
|
||||
items, titles = SearchAll.collect_all()
|
||||
results = SearchAll.get_results(titles, query)
|
||||
results = [items[i[2]] for i in results]
|
||||
def get_track_items(item: dict[str, _type], query: str, limit=5):
|
||||
tracks: list[Track] = []
|
||||
|
||||
return SearchAll.sort_results(results)
|
||||
if item["type"] == "track":
|
||||
tracks.extend(SearchTracks(query)())
|
||||
|
||||
if item["type"] == "album":
|
||||
t = TrackStore.get_tracks_by_albumhash(item["item"].albumhash)
|
||||
t.sort(key=lambda x: x.last_mod)
|
||||
|
||||
# if there are less than the limit, get more tracks
|
||||
if len(t) < limit:
|
||||
remainder = limit - len(t)
|
||||
more_tracks = SearchTracks(query)()
|
||||
t.extend(more_tracks[:remainder])
|
||||
|
||||
tracks.extend(t)
|
||||
|
||||
if item["type"] == "artist":
|
||||
t = TrackStore.get_tracks_by_artisthash(item["item"].artisthash)
|
||||
t.sort(key=lambda x: x.last_mod)
|
||||
|
||||
# if there are less than the limit, get more tracks
|
||||
if len(t) < limit:
|
||||
remainder = limit - len(t)
|
||||
more_tracks = SearchTracks(query)()
|
||||
t.extend(more_tracks[:remainder])
|
||||
|
||||
tracks.extend(t)
|
||||
|
||||
return tracks[:limit]
|
||||
|
||||
@staticmethod
|
||||
def get_album_items(item: dict[str, _type], query: str, limit=6):
|
||||
if item["type"] == "track":
|
||||
return SearchAlbums(query)()[:limit]
|
||||
|
||||
if item["type"] == "album":
|
||||
return SearchAlbums(query)()[:limit]
|
||||
|
||||
if item["type"] == "artist":
|
||||
albums = AlbumStore.get_albums_by_artisthash(item["item"].artisthash)
|
||||
|
||||
# if there are less than the limit, get more albums
|
||||
if len(albums) < limit:
|
||||
remainder = limit - len(albums)
|
||||
more_albums = SearchAlbums(query)()
|
||||
albums.extend(more_albums[:remainder])
|
||||
|
||||
return albums[:limit]
|
||||
|
||||
@staticmethod
|
||||
def search(query: str, albums_only=False, tracks_only=False, in_quotes=False):
|
||||
items, titles = TopResults.collect_all()
|
||||
results = TopResults.get_results(titles, query)
|
||||
|
||||
tracks_limit = Limit.tracks if tracks_only else 5
|
||||
albums_limit = Limit.albums if albums_only else 6
|
||||
artists_limit = 3
|
||||
|
||||
# map results to their respective items
|
||||
try:
|
||||
result = [items[i[2]] for i in results][0]
|
||||
except IndexError:
|
||||
if tracks_only:
|
||||
return []
|
||||
|
||||
if albums_only:
|
||||
return []
|
||||
|
||||
return {
|
||||
"top_result": None,
|
||||
"tracks": [],
|
||||
"artists": [],
|
||||
"albums": [],
|
||||
}
|
||||
|
||||
result = TopResults.map_with_type(result)
|
||||
|
||||
if in_quotes:
|
||||
top_tracks = SearchTracks(query)()[:tracks_limit]
|
||||
else:
|
||||
top_tracks = TopResults.get_track_items(result, query, limit=tracks_limit)
|
||||
|
||||
if tracks_only:
|
||||
return top_tracks
|
||||
|
||||
if in_quotes:
|
||||
albums = SearchAlbums(query)()[:albums_limit]
|
||||
else:
|
||||
albums = TopResults.get_album_items(result, query, limit=albums_limit)
|
||||
|
||||
if albums_only:
|
||||
return albums
|
||||
|
||||
artists = SearchArtists(query)()[:artists_limit]
|
||||
|
||||
return {
|
||||
"top_result": result,
|
||||
"tracks": top_tracks,
|
||||
"artists": artists,
|
||||
"albums": albums,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user