sublime-music/sublime_music/adapters/api_objects.py

"""
Defines the objects that are returned by adapter methods.
"""
import abc
import logging
from datetime import datetime, timedelta
from functools import lru_cache, partial
from enum import Enum
from typing import (
    Any,
    Callable,
    cast,
    Dict,
    Iterable,
    List,
    Optional,
    Sequence,
    Tuple,
    TypeVar,
    Union,
)

from fuzzywuzzy import fuzz


class Genre(abc.ABC):
    name: str
    song_count: Optional[int]
    album_count: Optional[int]


class Album(abc.ABC):
    """
    The ``id`` field is optional, because there are some situations where an adapter
    (such as Subsonic) sends an album name, but not an album ID.
    """

    name: str
    id: Optional[str]
    artist: Optional["Artist"]
    cover_art: Optional[str]
    created: Optional[datetime]
    duration: Optional[timedelta]
    genre: Optional[Genre]
    play_count: Optional[int]
    song_count: Optional[int]
    songs: Optional[Sequence["Song"]]
    starred: Optional[datetime]
    year: Optional[int]


class Artist(abc.ABC):
    """
    The ``id`` field is optional, because there are some situations where an adapter
    (such as Subsonic) sends an artist name, but not an artist ID. This especially
    happens when there are multiple artists.
    """

    name: str
    id: Optional[str]
    album_count: Optional[int]
    artist_image_url: Optional[str]
    starred: Optional[datetime]
    albums: Optional[Sequence[Album]]

    similar_artists: Optional[Sequence["Artist"]] = None
    biography: Optional[str] = None
    music_brainz_id: Optional[str] = None
    last_fm_url: Optional[str] = None


class Directory(abc.ABC):
    """
    The special directory with ``name`` and ``id`` should be used to indicate the
    top-level directory.
    """

    id: str
    name: Optional[str]
    parent_id: Optional[str]
    children: Sequence[Union["Directory", "Song"]]


class Song(abc.ABC):
    id: str
    title: str
    path: Optional[str]
    parent_id: Optional[str]
    duration: Optional[timedelta]

    album: Optional[Album]
    artist: Optional[Artist]
    genre: Optional[Genre]

    track: Optional[int]
    disc_number: Optional[int]
    year: Optional[int]
    cover_art: Optional[str]
    size: Optional[int]
    user_rating: Optional[int]
    starred: Optional[datetime]


class Playlist(abc.ABC):
    id: str
    name: str
    song_count: Optional[int]
    duration: Optional[timedelta]
    songs: Sequence[Song]
    created: Optional[datetime]
    changed: Optional[datetime]
    comment: Optional[str]
    owner: Optional[str]
    public: Optional[bool]
    cover_art: Optional[str]


class PlayQueue(abc.ABC):
    songs: Sequence[Song]
    position: timedelta
    username: Optional[str]
    changed: Optional[datetime]
    changed_by: Optional[str]
    value: Optional[str]
    current_index: Optional[int]


@lru_cache(maxsize=8192)
def similarity_ratio(query: str, string: str) -> int:
    """
    Return the :class:`fuzzywuzzy.fuzz.partial_ratio` between the ``query`` and
    the given ``string``.

    This ends up being called quite a lot, so the result is cached in an LRU
    cache using :class:`functools.lru_cache`.

    :param query: the query string
    :param string: the string to compare to the query string
    """
    return fuzz.partial_ratio(query, string)


class SearchResult:
    """
    An object representing the aggregate results of a search which can include
    both server and local results.
    """

    class Kind(Enum):
        ARTIST = 0
        ALBUM = 1
        SONG = 2
        PLAYLIST = 2

    ValueType = Union[Artist, Album, Song, Playlist]

    def __init__(self, query: str = None):
        self.query = query
        self.similiarity_partial = partial(
            similarity_ratio, self.query.lower() if self.query else ""
        )
        self._results: Dict[Tuple[Kind, str], ValueType] = {}

    def __repr__(self) -> str:
        fields = ("query", "_results")
        formatted_fields = map(lambda f: f"{f}={getattr(self, f)}", fields)
        return f"<SearchResult {' '.join(formatted_fields)}>"

    def add_results(self, result_type: str, results: Iterable):
        """Adds the ``results`` to the ``_result_type`` set."""
        if results is None:
            return

        for result in results:
            if result_type == 'artists':
                kind = self.Kind.ARTIST
            elif result_type == 'albums':
                kind = self.Kind.ALBUM
            elif result_type == 'songs':
                kind = self.Kind.SONG
            elif result_type == 'playlists':
                kind = self.Kind.PLAYLIST
            else:
                assert False

            self._results[(kind, result.id)] = result

    def update(self, other: "SearchResult"):
        assert self.query == other.query
        self._results.update(other._results)

    def _transform(self, kind: Kind, value: ValueType) -> Tuple[str, ...]:
        if kind is self.Kind.ARTIST:
            return (value.name,)
        elif kind is self.Kind.ALBUM:
            return (value.name, value.artist and value.artist.name)
        elif kind is self.Kind.SONG:
            return (value.title, value.artist and value.artist.name)
        elif kind is self.Kind.PLAYLIST:
            return (value.name,)
        else:
            assert False

    def get_results(self) -> List[Tuple[Kind, ValueType]]:
        assert self.query

        all_results = []
        for (kind, _), value in self._results.items():
            try:
                transformed = self._transform(kind, value)
            except Exception:
                continue

            max_similarity = max(
                (self.similiarity_partial(t.lower()) for t in transformed
                                                     if t is not None),
                default=0)

            if max_similarity < 60:
                continue

            all_results.append((max_similarity, (kind, value)))

        all_results.sort(key=lambda rx: rx[0], reverse=True)

        logging.debug(similarity_ratio.cache_info())
        return [r for _, r in all_results]