beets/beetsplug/edit.py

"""Open metadata information in a text editor to let the user edit it."""

from __future__ import annotations

import codecs
import os
import shlex
import subprocess
from collections import Counter
from tempfile import NamedTemporaryFile
from typing import TYPE_CHECKING, Any, cast

import yaml

from beets import plugins, ui, util
from beets.dbcore import types
from beets.exceptions import UserError
from beets.importer import Action
from beets.library import Album, Item
from beets.ui.commands.utils import do_query
from beets.util import PromptChoice

if TYPE_CHECKING:
    from beets.importer import ImportSession, ImportTask

# These "safe" types can avoid the format/parse cycle that most fields go
# through: they are safe to edit with native YAML types.
SAFE_TYPES = (
    types.BaseFloat,
    types.BaseInteger,
    types.Boolean,
    types.DelimitedString,
)

# Fixed fields that only exist on Item, not Album (e.g. title, track, path).
# Flexible attributes are deliberately excluded from this set: they aren't
# part of either model's fixed schema, so they can't be told apart by name
# alone and are left for the user to configure correctly.
ITEM_ONLY_FIELDS = Item._field_names - Album._field_names


class ParseError(Exception):
    """The modified file is unreadable. The user should be offered a chance to
    fix the error.
    """


def edit(filename, log):
    """Open `filename` in a text editor."""
    cmd = shlex.split(util.editor_command())
    cmd.append(filename)
    log.debug("invoking editor command: {!r}", cmd)
    try:
        subprocess.call(cmd)
    except OSError as exc:
        raise UserError(f"could not run editor command {cmd[0]!r}: {exc}")


def dump(arg):
    """Dump a sequence of dictionaries as YAML for editing."""
    return yaml.safe_dump_all(arg, allow_unicode=True, default_flow_style=False)


def load(s):
    """Read a sequence of YAML documents back to a list of dictionaries
    with string keys.

    Can raise a `ParseError`.
    """
    try:
        out = []
        for d in yaml.safe_load_all(s):
            if not isinstance(d, dict):
                raise ParseError(
                    f"each entry must be a dictionary; found {type(d).__name__}"
                )

            # Convert all keys to strings. They started out as strings,
            # but the user may have inadvertently messed this up.
            out.append({str(k): v for k, v in d.items()})

    except yaml.YAMLError as e:
        raise ParseError(f"invalid YAML: {e}")
    return out


def _safe_value(obj, key, value):
    """Check whether the `value` is safe to represent in YAML and trust as
    returned from parsed YAML.

    This ensures that values do not change their type when the user edits their
    YAML representation.
    """
    if value is None:
        return True
    typ = obj._type(key)
    return isinstance(typ, SAFE_TYPES) and isinstance(value, typ.model_type)


def flatten(obj, fields):
    """Represent `obj`, a `dbcore.Model` object, as a dictionary for
    serialization. Only include the given `fields` if provided;
    otherwise, include everything.

    The resulting dictionary's keys are strings and the values are
    safely YAML-serializable types.
    """
    # Format each value.
    d = {}
    for key in obj.keys():
        value = obj[key]
        if _safe_value(obj, key, value):
            # A safe value that is faithfully representable in YAML.
            d[key] = value
        else:
            # A value that should be edited as a string.
            d[key] = obj.formatted()[key]

    # Possibly filter field names.
    if fields:
        return {k: v for k, v in d.items() if k in fields}
    return d


def apply_(obj, data):
    """Set the fields of a `dbcore.Model` object according to a
    dictionary.

    This is the opposite of `flatten`. The `data` dictionary should have
    strings as values.
    """
    for key, value in data.items():
        if _safe_value(obj, key, value):
            # A safe value *stayed* represented as a safe type. Assign it
            # directly.
            obj[key] = value
        else:
            # Either the field was stringified originally or the user changed
            # it from a safe type to an unsafe one. Parse it as a string.
            obj.set_parse(key, str(value))


class EditPlugin(plugins.BeetsPlugin):
    def __init__(self):
        super().__init__()

        self.config.add(
            {
                # The default fields to edit.
                "albumfields": "album albumartist",
                "itemfields": "track title artist album",
                # Silently ignore any changes to these fields.
                "ignore_fields": "id path",
            }
        )

        self.register_listener(
            "before_choose_candidate", self.before_choose_candidate_listener
        )

    def commands(self):
        edit_command = ui.Subcommand("edit", help="interactively edit metadata")
        edit_command.parser.add_option(
            "-f",
            "--field",
            metavar="FIELD",
            action="append",
            help="edit this field also",
        )
        edit_command.parser.add_option(
            "--all", action="store_true", dest="all", help="edit all fields"
        )
        edit_command.parser.add_album_option()
        edit_command.func = self._edit_command
        return [edit_command]

    def _edit_command(self, lib, opts, args):
        """The CLI command function for the `beet edit` command."""
        # Get the objects to edit.
        items, albums = do_query(lib, args, opts.album, False)
        objs = albums if opts.album else items
        if not objs:
            ui.print_("Nothing to edit.")
            return

        # Get the fields to edit.
        if opts.all:
            fields = None
        else:
            fields = self._get_fields(opts.album, opts.field)
        self.edit(opts.album, objs, fields)

    def _get_fields(self, album, extra):
        """Get the set of fields to edit."""
        # Start with the configured base fields.
        if album:
            fields = self.config["albumfields"].as_str_seq()
        else:
            fields = self.config["itemfields"].as_str_seq()

        # Add the requested extra fields.
        if extra:
            fields += extra

        # Ensure we always have the `id` field for identification.
        fields.append("id")

        return set(fields)

    def edit(self, album, objs, fields):
        """The core editor function.

        - `album`: A flag indicating whether we're editing Items or Albums.
        - `objs`: The `Item`s or `Album`s to edit.
        - `fields`: The set of field names to edit (or None to edit
          everything).
        """
        # Present the YAML to the user and let them change it.
        success = self.edit_objects(objs, fields)

        # Save the new data.
        if success:
            self.save_changes(objs)

    def edit_objects(self, objs, fields):
        """Dump a set of Model objects to a file as text, ask the user
        to edit it, and apply any changes to the objects.

        Return a boolean indicating whether the edit succeeded.
        """
        # Get the content to edit as raw data structures.
        old_data = [flatten(o, fields) for o in objs]
        cur_str = dump(old_data)

        # Loop until we have parseable data and the user confirms.
        while True:
            result = self._edit_yaml(cur_str)
            if result is None:
                return False
            new_data, new_str = result

            # Show the changes.
            # If the objects are not on the DB yet, we need a copy of their
            # original state for show_model_changes.
            objs_old = [obj.copy() if obj.id < 0 else None for obj in objs]
            self.apply_data(objs, old_data, new_data)
            changed = False
            for obj, obj_old in zip(objs, objs_old):
                changed |= ui.show_model_changes(obj, obj_old)
            if not changed:
                ui.print_("No changes to apply.")
                return False

            # For cancel/keep-editing, restore objects to their original
            # in-memory state so temp edits don't leak into the session
            choice = ui.input_options(("continue Editing", "apply", "cancel"))
            if choice == "a":  # Apply.
                return True
            if choice == "c":  # Cancel.
                self.apply_data(objs, new_data, old_data)
                return False
            if choice == "e":  # Keep editing.
                self.apply_data(objs, new_data, old_data)
                cur_str = new_str
                continue

    def apply_data(self, objs, old_data, new_data):
        """Take potentially-updated data and apply it to a set of Model
        objects.

        Documents are matched to objects by their ``id`` field rather than
        by position, so a reordered or otherwise misaligned document list
        cannot cause one object's data to be applied to a different object.

        The objects are not written back to the database, so the changes
        are temporary.
        """
        if len(old_data) != len(new_data):
            self._log.warning(
                "number of objects changed from {} to {}",
                len(old_data),
                len(new_data),
            )

        obj_by_id = {o.id: o for o in objs}
        old_by_id = {d.get("id"): d for d in old_data}
        new_id_counts = Counter(d.get("id") for d in new_data)
        ignore_fields = self.config["ignore_fields"].as_str_seq()
        for new_dict in new_data:
            new_id = new_dict.get("id")
            if new_id_counts[new_id] > 1:
                # Two or more documents claim the same id: either a document's
                # id was edited to collide with another one, or the same
                # document was duplicated. We can't tell which document is
                # the "real" one, so ignore all of them.
                self._log.warning(
                    "ignoring objects with duplicate id {}", new_id
                )
                continue

            old_dict = old_by_id.get(new_id)
            obj = obj_by_id.get(new_id)
            if old_dict is None or obj is None:
                self._log.warning("ignoring object whose id changed")
                continue

            # Prohibit any changes to forbidden fields to avoid
            # clobbering `id` and such by mistake.
            forbidden = False
            for key in ignore_fields:
                if old_dict.get(key) != new_dict.get(key):
                    self._log.warning("ignoring object whose {} changed", key)
                    forbidden = True
                    break
            if forbidden:
                continue

            apply_(obj, new_dict)

    def save_changes(self, objs):
        """Save a list of updated Model objects to the database."""
        # Save to the database and possibly write tags.
        for ob in objs:
            if ob._dirty:
                self._log.debug("saving changes to {}", ob)
                ob.try_sync(ui.should_write(), ui.should_move())

    # Methods for interactive importer execution.

    def before_choose_candidate_listener(self, session, task):
        """Append an "Edit" choice and an "edit Candidates" choice (if
        there are candidates) to the interactive importer prompt.
        """
        choices = [PromptChoice("d", "eDit", self.importer_edit)]
        if task.candidates:
            choices.append(
                PromptChoice(
                    "c", "edit Candidates", self.importer_edit_candidate
                )
            )

        return choices

    def _importer_edit_album_header(
        self, task: ImportTask
    ) -> dict[str, Any] | None:
        """Build the album-header YAML document for import editing.

        Returns a dict of album-level fields, or ``None`` when the current
        task is not an album import.
        """
        if not getattr(task, "is_album", False) or not task.items:
            return None

        album_fields = set(self.config["albumfields"].as_str_seq())
        if not album_fields:
            return None

        # Drop fields that only exist on Item (title, track, path, ...);
        # since the header is built from a single item and then applied to
        # every item, an item-only field here would silently stamp that
        # one item's value onto the whole album. Flexible fields are left
        # alone so they can still be edited at the album level.
        item_only_fields = album_fields & ITEM_ONLY_FIELDS
        if item_only_fields:
            self._log.warning(
                "ignoring item-only fields configured in albumfields: {}",
                ", ".join(sorted(item_only_fields)),
            )
            album_fields -= ITEM_ONLY_FIELDS
        if not album_fields:
            return None

        first_item = task.items[0]
        header = flatten(first_item, album_fields)
        return header if header else None

    def _importer_edit_apply_header(
        self, items: list[Item], header_data: dict[str, Any]
    ) -> None:
        """Apply album-header changes to every item in the list."""
        if not header_data:
            return
        for item in items:
            apply_(item, header_data)

    def _edit_yaml(
        self, old_str: str
    ) -> tuple[list[dict[str, Any]], str] | None:
        """Open a temporary file with `old_str`, let the user edit it, and
        return the parsed list of YAML documents.

        Returns ``(parsed_data, edited_str)`` on success, or ``None`` if the
        user aborted (no changes or unresolvable parse error).
        """
        with NamedTemporaryFile(
            mode="w", suffix=".yaml", delete=False, encoding="utf-8"
        ) as new:
            new.write(old_str)

        try:
            while True:
                edit(new.name, self._log)

                with codecs.open(new.name, encoding="utf-8") as f:
                    new_str = f.read()
                if new_str == old_str:
                    ui.print_("No changes; aborting.")
                    return None

                try:
                    return load(new_str), new_str
                except ParseError as e:
                    ui.print_(f"Could not read data: {e}")
                    if not ui.input_yn("Edit again to fix? (Y/n)", True):
                        return None
        finally:
            os.remove(new.name)

    def importer_edit(
        self, session: ImportSession, task: ImportTask
    ) -> Action | None:
        """Callback for invoking the functionality during an interactive
        import session on the *original* item tags.
        """
        # Assign negative temporary ids to Items that are not in the database
        # yet. By using negative values, no clash with items in the database
        # can occur.
        for i, obj in enumerate(task.items, start=1):
            # The importer may set the id to None when re-importing albums.
            if not obj._db or obj.id is None:
                obj.id = -i

        # Decide which fields to show.
        album_fields = set()
        if getattr(task, "is_album", False):
            album_fields = set(self.config["albumfields"].as_str_seq())
        item_fields = set(self.config["itemfields"].as_str_seq())

        # Track-level fields exclude any that are shown in the album header
        # to avoid duplication.
        track_fields = item_fields - album_fields
        track_fields.add("id")

        # Build the YAML document list.
        header_data = self._importer_edit_album_header(task)
        old_track_data = [flatten(o, track_fields) for o in task.items]
        all_old_data = []
        if header_data is not None:
            all_old_data.append(header_data)
        all_old_data.extend(old_track_data)
        has_header = header_data is not None
        num_data_docs = 1 if has_header else 0

        cur_str = dump(all_old_data)

        while True:
            result = self._edit_yaml(cur_str)
            if result is None:
                self._importer_edit_cleanup(task)
                return None
            new_all_data, new_str = result

            expected_total = num_data_docs + len(task.items)
            if len(new_all_data) != expected_total:
                ui.print_(
                    f"Number of documents changed from {expected_total} to "
                    f"{len(new_all_data)}."
                )
                if ui.input_yn("Edit again to fix? (Y/n)", True):
                    continue
                self._importer_edit_cleanup(task)
                return None

            # Split into album header and per-track documents. Every track
            # document always has an `id` field (see `track_fields` above),
            # while the header never does, so identify the header by the
            # absence of `id` rather than by position. This keeps the split
            # correct even if the user moves the header elsewhere in the
            # file, since `apply_data` already matches track documents by
            # id regardless of order.
            if has_header:
                new_header_data = [d for d in new_all_data if "id" not in d]
                new_track_data = [d for d in new_all_data if "id" in d]
                if len(new_header_data) != 1:
                    ui.print_(
                        "Could not identify the album header: exactly one "
                        "document must have no `id` field."
                    )
                    if ui.input_yn("Edit again to fix? (Y/n)", True):
                        continue
                    self._importer_edit_cleanup(task)
                    return None
            else:
                new_header_data = []
                new_track_data = new_all_data

            # Snapshot originals for diff display and restore.
            objs_old = cast("list[Item]", [obj.copy() for obj in task.items])

            # Apply header changes to every item.
            if new_header_data:
                self._importer_edit_apply_header(task.items, new_header_data[0])

            # Apply per-track changes.
            self.apply_data(task.items, old_track_data, new_track_data)

            # Show the diff.
            changed = False
            for item, old_copy in zip(task.items, objs_old):
                changed |= ui.show_model_changes(item, old_copy)

            if not changed:
                ui.print_("No changes to apply.")
                self._importer_edit_cleanup(task)
                return None

            choice = ui.input_options(("continue Editing", "apply", "cancel"))
            if choice == "a":  # Apply.
                self._importer_edit_cleanup(task)
                return Action.RETAG
            if choice == "c":  # Cancel.
                self._importer_edit_restore_from_copies(task, objs_old)
                self._importer_edit_cleanup(task)
                return None
            if choice == "e":  # Keep editing.
                self._importer_edit_restore_from_copies(task, objs_old)
                cur_str = new_str
                continue

    @staticmethod
    def _importer_edit_cleanup(task: ImportTask) -> None:
        """Remove temporary negative ids from task items."""
        for obj in task.items:
            if obj.id is not None and obj.id < 0:
                obj.id = None

    @staticmethod
    def _importer_edit_restore_from_copies(
        task: ImportTask, copies: list[Item]
    ) -> None:
        """Restore items to their state before the last edit cycle.

        ``copies`` must be a list of :class:`Item <beets.library.Item>` copies
        taken *before* the changes were applied.
        """
        for i, item in enumerate(task.items):
            if i < len(copies):
                for key in item._fields:
                    item[key] = copies[i][key]

    def importer_edit_candidate(self, session, task):
        """Callback for invoking the functionality during an interactive
        import session on a *candidate*. The candidate's metadata is
        applied to the original items.
        """
        # Prompt the user for a candidate.
        sel = ui.input_options([], numrange=(1, len(task.candidates)))
        # Force applying the candidate on the items.
        task.match = task.candidates[sel - 1]
        task.apply_metadata()

        return self.importer_edit(session, task)