Skip to main content
`Ctrl`+`K`

# Source code for beets.metadata_plugins

    """Metadata source plugin interface.

    This allows beets to lookup metadata from various sources. We define
    a common interface for all metadata sources which need to be
    implemented as plugins.
    """

    from __future__ import annotations

    import abc
    import re
    from concurrent.futures import ThreadPoolExecutor, as_completed
    from contextlib import contextmanager
    from functools import cache, cached_property, wraps
    from typing import (
        TYPE_CHECKING,
        Generic,
        Literal,
        NamedTuple,
        TypedDict,
        TypeVar,
    )

    import unidecode
    from confuse import NotFoundError

    from beets import config, logging
    from beets.util import cached_classproperty
    from beets.util.id_extractors import extract_release_id

    from .plugins import BeetsPlugin, find_plugins, notify_info_yielded, send

    Ret = TypeVar("Ret")
    QueryType = Literal["album", "track"]

    if TYPE_CHECKING:
        from collections.abc import Callable, Iterable, Iterator, Sequence

        from .autotag import AlbumInfo, TrackInfo
        from .library.models import Item

    # Global logger.
    log = logging.getLogger("beets")

    @cache
    def find_metadata_source_plugins() -> list[MetadataSourcePlugin]:
        """Return a list of all loaded metadata source plugins."""
        # TODO: Make this an isinstance(MetadataSourcePlugin, ...) check in v3.0.0
        # This should also allow us to remove the type: ignore comments below.
        return [p for p in find_plugins() if hasattr(p, "data_source")]  # type: ignore[misc]

    @cache
    def get_metadata_source(name: str) -> MetadataSourcePlugin | None:
        """Get metadata source plugin by name."""
        name = name.lower()
        plugins = find_metadata_source_plugins()
        return next((p for p in plugins if p.data_source.lower() == name), None)

    @contextmanager
    def maybe_handle_plugin_error(plugin: MetadataSourcePlugin, method_name: str):
        """Safely call a plugin method, catching and logging exceptions."""
        if config["raise_on_error"]:
            yield
        else:
            try:
                yield
            except Exception as e:
                log.error(
                    "Error in '{}.{}': {}", plugin.data_source, method_name, e
                )
                log.debug("Exception details:", exc_info=True)

    def _yield_from_plugins(
        func: Callable[..., Iterable[Ret]],
    ) -> Callable[..., Iterator[Ret]]:
        method_name = func.__name__

        def materialize(
            plugin: MetadataSourcePlugin, method_name: str, *args, **kwargs
        ) -> list[Ret]:
            method = getattr(plugin, method_name)
            return list(method(*args, **kwargs))

        @wraps(func)
        def wrapper(*args, **kwargs) -> Iterator[Ret]:
            # Run plugin methods concurrently for faster I/O-bound lookups.
            with ThreadPoolExecutor() as executor:
                futures = {
                    executor.submit(
                        # Evaluate iterator with list such that results are ready when
                        # future.result() is called.
                        materialize,
                        plugin,
                        method_name,
                        *args,
                        **kwargs,
                    ): plugin
                    for plugin in find_metadata_source_plugins()
                }

                for future in as_completed(futures):
                    plugin = futures[future]
                    with maybe_handle_plugin_error(plugin, method_name):
                        yield from filter(None, future.result())

        return wrapper

    @notify_info_yielded("albuminfo_received")
    @_yield_from_plugins
    def candidates(*args, **kwargs) -> Iterator[AlbumInfo]:
        yield from ()

    @notify_info_yielded("trackinfo_received")
    @_yield_from_plugins
    def item_candidates(*args, **kwargs) -> Iterator[TrackInfo]:
        yield from ()

    @notify_info_yielded("albuminfo_received")
    @_yield_from_plugins
    def albums_for_ids(*args, **kwargs) -> Iterator[AlbumInfo]:
        yield from ()

    @notify_info_yielded("trackinfo_received")
    @_yield_from_plugins
    def tracks_for_ids(*args, **kwargs) -> Iterator[TrackInfo]:
        yield from ()

    def album_for_id(_id: str, data_source: str) -> AlbumInfo | None:
        """Get AlbumInfo object for the given ID and data source."""
        if plugin := get_metadata_source(data_source):
            with maybe_handle_plugin_error(plugin, "album_for_id"):
                if info := plugin.album_for_id(_id):
                    send("albuminfo_received", info=info)
                    return info

        return None

    def track_for_id(_id: str, data_source: str) -> TrackInfo | None:
        """Get TrackInfo object for the given ID and data source."""
        if plugin := get_metadata_source(data_source):
            with maybe_handle_plugin_error(plugin, "track_for_id"):
                if info := plugin.track_for_id(_id):
                    send("trackinfo_received", info=info)
                    return info

        return None

    @cache
    def get_penalty(data_source: str | None) -> float:
        """Get the penalty value for the given data source."""
        return next(
            (
                p.data_source_mismatch_penalty
                for p in find_metadata_source_plugins()
                if p.data_source == data_source
            ),
            MetadataSourcePlugin.DEFAULT_DATA_SOURCE_MISMATCH_PENALTY,
        )

    class MetadataSourcePlugin(BeetsPlugin, metaclass=abc.ABCMeta):
        """A plugin that provides metadata from a specific source.

        This base class implements a contract for plugins that provide metadata
        from a specific source. The plugin must implement the methods to search for albums
        and tracks, and to retrieve album and track information by ID.
        """

        DEFAULT_DATA_SOURCE_MISMATCH_PENALTY = 0.5

        @cached_classproperty
        def data_source(cls) -> str:
            """The data source name for this plugin.

            This is inferred from the plugin name.
            """
            return cls.__name__.replace("Plugin", "")  # type: ignore[attr-defined]

        @cached_property
        def data_source_mismatch_penalty(self) -> float:
            try:
                return self.config["source_weight"].as_number()
            except NotFoundError:
                return self.config["data_source_mismatch_penalty"].as_number()

        def __init__(self, *args, **kwargs) -> None:
            super().__init__(*args, **kwargs)
            self.config.add(
                {
                    "search_limit": 5,
                    "data_source_mismatch_penalty": self.DEFAULT_DATA_SOURCE_MISMATCH_PENALTY,  # noqa: E501
                }
            )

        @abc.abstractmethod
        def album_for_id(self, album_id: str) -> AlbumInfo | None:
            """Return :py:class:`AlbumInfo` object or None if no matching release was
            found."""
            raise NotImplementedError

        @abc.abstractmethod
        def track_for_id(self, track_id: str) -> TrackInfo | None:
            """Return a :py:class:`TrackInfo` object or None if no matching release was
            found.
            """
            raise NotImplementedError

        # ---------------------------------- search ---------------------------------- #

        @abc.abstractmethod
        def candidates(
            self, items: Sequence[Item], artist: str, album: str, va_likely: bool
        ) -> Iterable[AlbumInfo]:
            """Return :py:class:`AlbumInfo` candidates that match the given album.

            Used in the autotag functionality to search for albums.

            :param items: List of items in the album
            :param artist: Album artist
            :param album: Album name
            :param va_likely: Whether the album is likely to be by various artists
            """
            raise NotImplementedError

        @abc.abstractmethod
        def item_candidates(
            self, item: Item, artist: str, title: str
        ) -> Iterable[TrackInfo]:
            """Return :py:class:`TrackInfo` candidates that match the given track.

            Used in the autotag functionality to search for tracks.

            :param item: Track item
            :param artist: Track artist
            :param title: Track title
            """
            raise NotImplementedError

        def albums_for_ids(self, ids: Iterable[str]) -> Iterable[AlbumInfo | None]:
            """Batch lookup of album metadata for a list of album IDs.

            Given a list of album identifiers, yields corresponding AlbumInfo objects.
            Missing albums result in None values in the output iterator.
            Plugins may implement this for optimized batched lookups instead of
            single calls to album_for_id.
            """

            return (self.album_for_id(id_) for id_ in ids)

        def tracks_for_ids(self, ids: Iterable[str]) -> Iterable[TrackInfo | None]:
            """Batch lookup of track metadata for a list of track IDs.

            Given a list of track identifiers, yields corresponding TrackInfo objects.
            Missing tracks result in None values in the output iterator.
            Plugins may implement this for optimized batched lookups instead of
            single calls to track_for_id.
            """

            return (self.track_for_id(id_) for id_ in ids)

        def _extract_id(self, url: str) -> str | None:
            """Extract an ID from a URL for this metadata source plugin.

            Uses the plugin's data source name to determine the ID format and
            extracts the ID from a given URL.
            """
            return extract_release_id(self.data_source, url)

        @staticmethod
        def get_artist(
            artists: Iterable[dict[str | int, str]],
            id_key: str | int = "id",
            name_key: str | int = "name",
            join_key: str | int | None = None,
        ) -> tuple[str, str | None]:
            """Returns an artist string (all artists) and an artist_id (the main
            artist) for a list of artist object dicts.

            For each artist, this function moves articles (such as 'a', 'an', and 'the')
            to the front. It returns a tuple containing the comma-separated string
            of all normalized artists and the ``id`` of the main/first artist.
            Alternatively a keyword can be used to combine artists together into a
            single string by passing the join_key argument.

            :param artists: Iterable of artist dicts or lists returned by API.
            :param id_key: Key or index corresponding to the value of ``id`` for
                the main/first artist. Defaults to 'id'.
            :param name_key: Key or index corresponding to values of names
                to concatenate for the artist string (containing all artists).
                Defaults to 'name'.
            :param join_key: Key or index corresponding to a field containing a
                keyword to use for combining artists into a single string, for
                example "Feat.", "Vs.", "And" or similar. The default is None
                which keeps the default behaviour (comma-separated).
            :return: Normalized artist string.
            """
            artist_id = None
            artist_string = ""
            artists = list(artists)  # In case a generator was passed.
            total = len(artists)
            for idx, artist in enumerate(artists):
                if not artist_id:
                    artist_id = artist[id_key]
                name = artist[name_key]
                # Move articles to the front.
                name = re.sub(r"^(.*?), (a|an|the)$", r"\2 \1", name, flags=re.I)
                # Use a join keyword if requested and available.
                if idx < (total - 1):  # Skip joining on last.
                    if join_key and artist.get(join_key, None):
                        name += f" {artist[join_key]} "
                    else:
                        name += ", "
                artist_string += name

            return artist_string, artist_id

    class IDResponse(TypedDict):
        """Response from the API containing an ID."""

        id: str

    class SearchParams(NamedTuple):
        """Bundle normalized search context passed to provider search hooks.

        Shared search orchestration constructs this value so plugin hooks receive
        one object describing search intent, query text, and provider filters.
        """

        query_type: QueryType
        query: str
        filters: dict[str, str]
        limit: int

    R = TypeVar("R", bound=IDResponse)

    class SearchApiMetadataSourcePlugin(
        Generic[R], MetadataSourcePlugin, metaclass=abc.ABCMeta
    ):
        """Helper class to implement a metadata source plugin with an API.

        Plugins using this ABC must implement an API search method to
        retrieve album and track information by ID,
        i.e. `album_for_id` and `track_for_id`, and a search method to
        perform a search on the API. The search method should return a list
        of identifiers for the requested type (album or track).
        """

        def __init__(self, *args, **kwargs) -> None:
            super().__init__(*args, **kwargs)
            self.config.add({"search_query_ascii": False})

        @abc.abstractmethod
        def get_search_query_with_filters(
            self,
            query_type: QueryType,
            items: Sequence[Item],
            artist: str,
            name: str,
            va_likely: bool,
        ) -> tuple[str, dict[str, str]]:
            """Build query text and API filters for a provider search.

            Subclasses can override this hook when their API requires a query format
            or filter set that differs from the default text-based construction.

            :param query_type: The type of query to perform. Either *album* or *track*
            :param items: List of items the search is being performed for
            :param artist: Artist name
            :param name: Album or track name, depending on ``query_type``
            :param va_likely: Whether the search is likely to be for various artists
            :return: Tuple of (``query`` text, ``filters`` dict) to use for the
                search API call
            """

        @abc.abstractmethod
        def get_search_response(self, params: SearchParams) -> Sequence[R]:
            """Fetch raw search results for a provider request.

            Implementations should return records containing source IDs so shared
            candidate resolution can perform ID-based album and track lookups.

            :param params: :py:namedtuple:`~SearchParams` named tuple
            :return: Sequence of IDResponse dicts containing at least an "id" key for each
            """

            raise NotImplementedError

        def _search_api(
            self, query_type: QueryType, query: str, filters: dict[str, str]
        ) -> Sequence[R]:
            """Run shared provider search orchestration and return ID-bearing results.

            This path applies optional query normalization and default limits, then
            delegates API access to provider hooks with consistent logging and
            failure handling.
            """
            if self.config["search_query_ascii"].get():
                query = unidecode.unidecode(query)

            limit = self.config["search_limit"].get(int)
            params = SearchParams(query_type, query, filters, limit)

            self._log.debug("Searching for '{}' with {}", query, filters)
            try:
                response_data = self.get_search_response(params)
            except Exception as e:
                if config["raise_on_error"].get(bool):
                    raise
                self._log.error(
                    "Error searching {.data_source}: {}", self, e, exc_info=True
                )
                return ()

            self._log.debug("Found {} result(s)", len(response_data))
            return response_data

        def _get_candidates(
            self, query_type: QueryType, *args, **kwargs
        ) -> Sequence[R]:
            """Resolve query hooks and execute one provider search request."""

            return self._search_api(
                query_type,
                *self.get_search_query_with_filters(query_type, *args, **kwargs),
            )

        def candidates(
            self, items: Sequence[Item], artist: str, album: str, va_likely: bool
        ) -> Iterable[AlbumInfo]:
            results = self._get_candidates("album", items, artist, album, va_likely)
            return filter(None, self.albums_for_ids(r["id"] for r in results))

        def item_candidates(
            self, item: Item, artist: str, title: str
        ) -> Iterable[TrackInfo]:
            results = self._get_candidates("track", [item], artist, title, False)
            return filter(None, self.tracks_for_ids(r["id"] for r in results))