beets/beetsplug/zero.py

"""Clears tag fields in media files."""

from __future__ import annotations

import re
from typing import TYPE_CHECKING, Any

import confuse
from mediafile import MediaFile

from beets.importer import Action
from beets.plugins import BeetsPlugin
from beets.ui import Subcommand, input_yn

if TYPE_CHECKING:
    import optparse

    from beets.importer import ImportSession, ImportTask
    from beets.library import Item, Library


__author__ = "baobab@heresiarch.info"


ARTWORK_FIELDS = {"images", "art"}


class ZeroPlugin(BeetsPlugin):
    fields_to_progs: dict[str, list[re.Pattern[str]]]

    def __init__(self) -> None:
        super().__init__()

        self.register_listener("write", self.write_event)
        self.register_listener(
            "import_task_choice", self.import_task_choice_event
        )

        self.config.add(
            {
                "auto": True,
                "fields": [],
                "keep_fields": [],
                "update_database": False,
                "omit_single_disc": False,
            }
        )

        self.fields_to_progs = {}
        self.warned = False

        """Read the bulk of the config into `self.fields_to_progs`.
        After construction, `fields_to_progs` contains all the fields that
        should be zeroed as keys and maps each of those to a list of compiled
        regexes (progs) as values.
        A field is zeroed if its value matches one of the associated progs. If
        progs is empty, then the associated field is always zeroed.
        """
        if self.config["fields"] and self.config["keep_fields"]:
            self._log.warning("cannot blacklist and whitelist at the same time")
        # Blacklist mode.
        elif self.config["fields"]:
            for field in self.config["fields"].as_str_seq():
                self._set_pattern(field)
        # Whitelist mode.
        elif self.config["keep_fields"]:
            keep = set(self.config["keep_fields"].as_str_seq())
            # ensure that all artwork fields are added when at least
            # one of them is present
            if keep & ARTWORK_FIELDS:
                keep.update(ARTWORK_FIELDS)
            for field in MediaFile.fields():
                if field not in keep and field not in (
                    "id",
                    "path",
                    "album_id",
                ):
                    self._set_pattern(field)

    def commands(self):
        zero_command = Subcommand("zero", help="set fields to null")

        def zero_fields(
            lib: Library, opts: optparse.Values, args: list[str]
        ) -> None:
            if not args and not input_yn(
                "Remove fields for all items? (Y/n)", True
            ):
                return
            for item in lib.items(args):
                self.process_item(item)

        zero_command.func = zero_fields
        return [zero_command]

    def _set_pattern(self, field):
        """Populate `self.fields_to_progs` for a given field.
        Do some sanity checks then compile the regexes.
        """
        if field not in MediaFile.fields():
            self._log.error("invalid field: {}", field)
        elif field in ("id", "path", "album_id"):
            self._log.warning(
                "field '{}' ignored, zeroing it would be dangerous", field
            )
        else:
            try:
                for pattern in self.config[field].as_str_seq():
                    prog = re.compile(pattern, re.IGNORECASE)
                    self.fields_to_progs.setdefault(field, []).append(prog)
            except confuse.NotFoundError:
                # Matches everything
                self.fields_to_progs[field] = []

    def import_task_choice_event(
        self, session: ImportSession, task: ImportTask
    ) -> None:
        if task.choice_flag == Action.ASIS and not self.warned:
            self._log.warning('cannot zero in "as-is" mode')
            self.warned = True
        # TODO request write in as-is mode

    def write_event(
        self, item: Item, path: bytes, tags: dict[str, Any]
    ) -> None:
        if self.config["auto"]:
            self.set_fields(item, tags)

    def set_fields(self, item, tags):
        """Set values in `tags` to `None` if the field is in
        `self.fields_to_progs` and any of the corresponding `progs` matches the
        field value.
        Also update the `item` itself if `update_database` is set in the
        config.
        """
        fields_set = False
        if self.config["omit_single_disc"].get(bool) and item.disctotal == 1:
            for tag in {"disc", "disctotal"} & set(tags):
                tags[tag] = None
                fields_set = True

        if not self.fields_to_progs:
            self._log.warning("no fields list to remove")

        for field, progs in self.fields_to_progs.items():
            if field in tags:
                value = tags[field]
                match = _match_progs(tags[field], progs)
            else:
                value = ""
                match = not progs

            if match:
                fields_set = True
                self._log.debug("{}: {} -> None", field, value)
                tags[field] = None
                if self.config["update_database"]:
                    item[field] = None

        return fields_set

    def process_item(self, item):
        tags = dict(item)

        if self.set_fields(item, tags):
            item.write(tags=tags)
            if self.config["update_database"]:
                item.store(fields=tags)


def _match_progs(value, progs):
    """Check if `value` (as string) is matching any of the compiled regexes in
    the `progs` list.
    """
    if not progs:
        return True
    for prog in progs:
        if prog.search(str(value)):
            return True
    return False