beets/beetsplug/missing.py

"""List missing tracks."""

from __future__ import annotations

from collections import defaultdict
from typing import TYPE_CHECKING, ClassVar

import requests

from beets import config, metadata_plugins
from beets.dbcore import types
from beets.library import Item
from beets.plugins import BeetsPlugin
from beets.ui import Subcommand, print_

from ._utils.musicbrainz import MusicBrainzAPIMixin

if TYPE_CHECKING:
    from collections.abc import Iterator

    from beets.library import Album, Library

# Valid MusicBrainz release types for filtering release groups
VALID_RELEASE_TYPES = [
    "nat",
    "album",
    "single",
    "ep",
    "broadcast",
    "other",
    "compilation",
    "soundtrack",
    "spokenword",
    "interview",
    "audiobook",
    "live",
    "remix",
    "dj-mix",
    "mixtape/street",
]

MB_ARTIST_QUERY = r"mb_albumartistid::^\w{8}-\w{4}-\w{4}-\w{4}-\w{12}$"


def _missing_count(album):
    """Return number of missing items in `album`."""
    return (album.albumtotal or 0) - len(album.items())


def _item(track_info, album_info, album_id):
    """Build and return `item` from `track_info` and `album info`
    objects. `item` is missing what fields cannot be obtained from
    MusicBrainz alone (encoder, rg_track_gain, rg_track_peak,
    rg_album_gain, rg_album_peak, original_year, original_month,
    original_day, length, bitrate, format, samplerate, bitdepth,
    channels, mtime.)
    """
    t = track_info
    a = album_info

    return Item(
        album_id=album_id,
        album=a.album,
        albumartist=a.artist,
        albumartist_credit=a.artist_credit,
        albumartist_sort=a.artist_sort,
        albumdisambig=a.albumdisambig,
        albumstatus=a.albumstatus,
        albumtype=a.albumtype,
        artist=t.artist,
        artist_credit=t.artist_credit,
        artist_sort=t.artist_sort,
        asin=a.asin,
        catalognum=a.catalognum,
        comp=a.va,
        country=a.country,
        day=a.day,
        disc=t.medium,
        disctitle=t.disctitle,
        disctotal=a.mediums,
        label=a.label,
        language=a.language,
        length=t.length,
        mb_albumid=a.album_id,
        mb_artistid=t.artist_id,
        mb_releasegroupid=a.releasegroup_id,
        mb_trackid=t.track_id,
        media=t.media,
        month=a.month,
        script=a.script,
        title=t.title,
        track=t.index,
        tracktotal=len(a.tracks),
        year=a.year,
    )


class MissingPlugin(MusicBrainzAPIMixin, BeetsPlugin):
    """List missing tracks"""

    album_types: ClassVar[dict[str, types.Type]] = {"missing": types.INTEGER}

    def __init__(self):
        super().__init__()

        self.config.add(
            {
                "count": False,
                "total": False,
                "album": False,
                "release_types": ["album"],
            }
        )

        self.album_template_fields["missing"] = _missing_count

        self._command = Subcommand("missing", help=__doc__, aliases=["miss"])
        self._command.parser.add_option(
            "-c",
            "--count",
            dest="count",
            action="store_true",
            help="count missing tracks per album",
        )
        self._command.parser.add_option(
            "-t",
            "--total",
            dest="total",
            action="store_true",
            help="count total of missing tracks",
        )
        self._command.parser.add_option(
            "-a",
            "--album",
            dest="album",
            action="store_true",
            help=(
                "show missing album releases for artist instead of tracks; "
                "Defaults to only releases of type 'album'"
            ),
        )
        self._command.parser.add_option(
            "--release-types",
            action="append",
            dest="release_types",
            help=(
                "comma-separated list of release types for missing albums "
                f"(valid: {', '.join(VALID_RELEASE_TYPES)})"
            ),
        )
        self._command.parser.add_format_option()

    def commands(self):
        def _miss(lib, opts, args):
            self.config.set_args(opts)
            albms = self.config["album"].get()

            helper = self._missing_albums if albms else self._missing_tracks
            helper(lib, args)

        self._command.func = _miss
        return [self._command]

    def _missing_tracks(self, lib, query):
        """Print a listing of tracks missing from each album in the library
        matching query.
        """
        albums = lib.albums(query)

        count = self.config["count"].get()
        total = self.config["total"].get()
        fmt = config["format_album" if count else "format_item"].get()

        if total:
            print(sum([_missing_count(a) for a in albums]))
            return

        # Default format string for count mode.
        if count:
            fmt += ": $missing"

        for album in albums:
            if count:
                if _missing_count(album):
                    print_(format(album, fmt))

            else:
                for item in self._missing(album):
                    print_(format(item, fmt))

    def _missing_albums(self, lib: Library, query: list[str]) -> None:
        """Print a listing of albums missing from each artist in the library
        matching query.
        """
        query.append(MB_ARTIST_QUERY)

        # build dict mapping artist to set of their release group ids in library
        album_ids_by_artist = defaultdict(set)
        for album in lib.albums(query):
            # TODO(@snejus): Some releases have different `albumartist` for the
            # same `mb_albumartistid`. Since we're grouping by the combination
            # of these two fields, we end up processing the same
            # `mb_albumartistid` multiple times: calling MusicBrainz API and
            # reporting the same set of missing albums. Instead, we should
            # group by `mb_albumartistid` field only.
            artist = (album["albumartist"], album["mb_albumartistid"])
            album_ids_by_artist[artist].add(album["mb_releasegroupid"])

        total_missing = 0
        release_types = []
        for rt in self.config["release_types"].as_str_seq():
            release_types.extend(rt.split(","))
        calculating_total = self.config["total"].get()
        for (artist, artist_id), album_ids in album_ids_by_artist.items():
            try:
                resp = self.mb_api.browse_release_groups(
                    artist=artist_id, type="|".join(release_types)
                )
            except requests.exceptions.RequestException:
                self._log.info(
                    "Couldn't fetch info for artist '{}' ({})",
                    artist,
                    artist_id,
                    exc_info=True,
                )
                continue

            missing_titles = [
                f"{artist} - {rg['title']}"
                for rg in resp
                if rg["id"] not in album_ids
            ]

            if calculating_total:
                total_missing += len(missing_titles)
            else:
                for title in missing_titles:
                    print(title)

        if calculating_total:
            print(total_missing)

    def _missing(self, album: Album) -> Iterator[Item]:
        """Query MusicBrainz to determine items missing from `album`."""
        if len(album.items()) == album.albumtotal:
            return

        # fetch missing items
        # TODO: Implement caching that without breaking other stuff
        data_source = album.get("data_source") or album.items()[0].get(
            "data_source", "MusicBrainz"
        )
        if album_info := metadata_plugins.album_for_id(
            album.mb_albumid, data_source
        ):
            item_mbids = {x.mb_trackid for x in album.items()}
            for track_info in album_info.tracks:
                if track_info.track_id not in item_mbids:
                    self._log.debug(
                        "track {.track_id} in album {.album_id}",
                        track_info,
                        album_info,
                    )
                    yield _item(track_info, album_info, album.id)