support in_quotes search query

This commit is contained in:
mungai-njoroge
2023-08-06 22:09:39 +03:00
parent 943d6e3590
commit f28d3f00bd
6 changed files with 238 additions and 174 deletions
+158 -46
View File
@@ -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,
}