""" This library contains all the functions related to the search functionality. """ from typing import List, Generator, TypeVar, Any import itertools from rapidfuzz import fuzz, process from unidecode import unidecode from app import models from app.utils.remove_duplicates import remove_duplicates from app.store.albums import AlbumStore from app.store.artists import ArtistStore from app.store.tracks import TrackStore ratio = fuzz.ratio wratio = fuzz.WRatio class Cutoff: """ Holds all the default cutoff values. """ tracks: int = 75 albums: int = 75 artists: int = 75 playlists: int = 75 class Limit: """ Holds all the default limit values. """ tracks: int = 150 albums: int = 150 artists: int = 150 playlists: int = 150 class SearchTracks: def __init__(self, query: str) -> None: self.query = query self.tracks = TrackStore.tracks def __call__(self) -> List[models.Track]: """ Gets all songs with a given title. """ track_titles = [unidecode(track.og_title).lower() for track in self.tracks] results = process.extract( self.query, track_titles, score_cutoff=Cutoff.tracks, limit=Limit.tracks, ) tracks = [self.tracks[i[2]] for i in results] return remove_duplicates(tracks) class SearchArtists: def __init__(self, query: str) -> None: self.query = query self.artists = ArtistStore.artists def __call__(self) -> list: """ Gets all artists with a given name. """ artists = [unidecode(a.name).lower() for a in self.artists] results = process.extract( self.query, artists, score_cutoff=Cutoff.artists, limit=Limit.artists, ) return [self.artists[i[2]] for i in results] class SearchAlbums: def __init__(self, query: str) -> None: self.query = query self.albums = AlbumStore.albums def __call__(self) -> List[models.Album]: """ Gets all albums with a given title. """ albums = [unidecode(a.og_title).lower() for a in self.albums] results = process.extract( self.query, albums, score_cutoff=Cutoff.albums, limit=Limit.albums, ) 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: self.playlists = playlists self.query = query def __call__(self) -> List[models.Playlist]: playlists = [p.name for p in self.playlists] results = process.extract( self.query, playlists, score_cutoff=Cutoff.playlists, limit=Limit.playlists, ) return [self.playlists[i[2]] for i in results] _type = List[models.Track | models.Album | models.Artist] _S2 = TypeVar("_S2") _ResultType = int | float def get_titles(items: _type): for item in items: if isinstance(item, models.Track): text = item.og_title elif isinstance(item, models.Album): text = item.title elif isinstance(item, models.Artist): text = item.name else: text = None yield text class SearchAll: """ Joins all tracks, albums and artists then fuzzy searches them as a single unit. """ @staticmethod def collect_all(): all_items: _type = [] all_items.extend(TrackStore.tracks) all_items.extend(AlbumStore.albums) all_items.extend(ArtistStore.artists) return all_items, get_titles(all_items) @staticmethod def get_results(items: Generator[str, Any, None], query: str): items = list(items) results = process.extract( query=query, choices=items, score_cutoff=Cutoff.tracks, limit=20 ) return results @staticmethod def sort_results(items: _type): """ Separates results into differrent lists using itertools.groupby. """ 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 ] mapped_items.sort(key=lambda x: x["type"]) groups = [ list(group) for key, group in itertools.groupby(mapped_items, lambda x: x["type"]) ] # 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 ] return groups @staticmethod def search(query: str): items, titles = SearchAll.collect_all() results = SearchAll.get_results(titles, query) results = [items[i[2]] for i in results] return SearchAll.sort_results(results)