beets/beetsplug/embedart.py

"""Allows beets to embed album art into file metadata."""

from __future__ import annotations

import os
import tempfile
from mimetypes import guess_extension
from typing import TYPE_CHECKING, Protocol

import requests

from beets import config, ui
from beets.exceptions import UserError
from beets.plugins import BeetsPlugin
from beets.ui import print_
from beets.util import bytestring_path, displayable_path, normpath, syspath
from beets.util.artresizer import ArtResizer
from beetsplug._utils import art

if TYPE_CHECKING:
    from collections.abc import Sequence

    from beets.importer import ImportSession, ImportTask
    from beets.library import Album, LibModel, Library


class EmbedArtCLIOpts(Protocol):
    file: str | None
    url: str | None
    yes: bool | None
    force: bool | None


class ExtractArtCLIOpts(Protocol):
    associate: bool | None
    filename: str | None
    outpath: str | None


class ClearArtCLIOpts(Protocol):
    yes: bool | None


def _confirm(objs: Sequence[LibModel], album: bool) -> bool:
    """Show the list of affected objects (items or albums) and confirm
    that the user wants to modify their artwork.

    `album` is a Boolean indicating whether these are albums (as opposed
    to items).
    """
    noun = "album" if album else "file"
    prompt = (
        f"Modify artwork for {len(objs)} {noun}{'s' if len(objs) > 1 else ''}?"
    )

    # Show all the items or albums.
    for obj in objs:
        print_(format(obj))

    # Confirm with user.
    return ui.input_yn(prompt)


class EmbedCoverArtPlugin(BeetsPlugin):
    """Allows albumart to be embedded into the actual files."""

    def __init__(self) -> None:
        super().__init__()
        self.config.add(
            {
                "maxwidth": 0,
                "auto": True,
                "compare_threshold": 0,
                "ifempty": False,
                "remove_art_file": False,
                "quality": 0,
                "clearart_on_import": False,
            }
        )

        if self.config["maxwidth"].get(int) and not ArtResizer.shared.local:
            self.config["maxwidth"] = 0
            self._log.warning(
                "ImageMagick or PIL not found; 'maxwidth' option ignored"
            )
        if (
            self.config["compare_threshold"].get(int)
            and not ArtResizer.shared.can_compare
        ):
            self.config["compare_threshold"] = 0
            self._log.warning(
                "ImageMagick 6.8.7 or higher not installed; "
                "'compare_threshold' option ignored"
            )

        self.register_listener("art_set", self.process_album)

        if self.config["clearart_on_import"].get(bool):
            self.register_listener("import_task_files", self.import_task_files)

    def commands(self) -> list[ui.Subcommand]:
        # Embed command.
        embed_cmd = ui.Subcommand(
            "embedart", help="embed image files into file metadata"
        )
        embed_cmd.parser.add_option(
            "-f", "--file", metavar="PATH", help="the image file to embed"
        )
        embed_cmd.parser.add_option(
            "--force", action="store_true", help="overwrite existing art"
        )

        embed_cmd.parser.add_option(
            "-y", "--yes", action="store_true", help="skip confirmation"
        )

        embed_cmd.parser.add_option(
            "-u",
            "--url",
            metavar="URL",
            help="the URL of the image file to embed",
        )

        maxwidth = self.config["maxwidth"].get(int)
        quality = self.config["quality"].get(int)
        compare_threshold = self.config["compare_threshold"].get(int)

        def embed_func(
            lib: Library, opts: EmbedArtCLIOpts, args: list[str]
        ) -> None:
            ifempty = not opts.force
            if opts.file:
                imagepath = normpath(opts.file)
                if not os.path.isfile(syspath(imagepath)):
                    raise UserError(
                        f"image file {displayable_path(imagepath)} not found"
                    )

                items = lib.items(args)

                # Confirm with user.
                if not opts.yes and not _confirm(items, not opts.file):
                    return

                for item in items:
                    art.embed_item(
                        self._log,
                        item,
                        imagepath,
                        maxwidth,
                        None,
                        compare_threshold,
                        ifempty,
                        quality=quality,
                    )
            elif opts.url:
                try:
                    response = requests.get(opts.url, timeout=5)
                    response.raise_for_status()
                except requests.exceptions.RequestException as e:
                    self._log.error("{}", e)
                    return
                extension = guess_extension(response.headers["Content-Type"])
                if extension is None:
                    self._log.error("Invalid image file")
                    return
                file = f"image{extension}"
                tempimg = os.fsencode(os.path.join(tempfile.gettempdir(), file))
                try:
                    with open(tempimg, "wb") as f:
                        f.write(response.content)
                except Exception as e:
                    self._log.error("Unable to save image: {}", e)
                    return
                items = lib.items(args)
                # Confirm with user.
                if not opts.yes and not _confirm(items, not opts.url):
                    os.remove(tempimg)
                    return
                for item in items:
                    art.embed_item(
                        self._log,
                        item,
                        tempimg,
                        maxwidth,
                        None,
                        compare_threshold,
                        ifempty,
                        quality=quality,
                    )
                os.remove(tempimg)
            else:
                albums = lib.albums(args)
                # Confirm with user.
                if not opts.yes and not _confirm(albums, not opts.file):
                    return
                for album in albums:
                    art.embed_album(
                        self._log,
                        album,
                        maxwidth,
                        False,
                        compare_threshold,
                        ifempty,
                        quality=quality,
                    )
                    self.remove_artfile(album)

        embed_cmd.func = embed_func

        # Extract command.
        extract_cmd = ui.Subcommand(
            "extractart", help="extract an image from file metadata"
        )
        extract_cmd.parser.add_option(
            "-o", dest="outpath", help="image output file"
        )
        extract_cmd.parser.add_option(
            "-n",
            dest="filename",
            help="image filename to create for all matched albums",
        )
        extract_cmd.parser.add_option(
            "-a",
            dest="associate",
            action="store_true",
            help="associate the extracted images with the album",
        )

        def extract_func(
            lib: Library, opts: ExtractArtCLIOpts, args: list[str]
        ) -> None:
            if opts.outpath:
                art.extract_first(
                    self._log, normpath(opts.outpath), lib.items(args)
                )
            else:
                filename = bytestring_path(
                    opts.filename or config["art_filename"].get()
                )
                if os.path.dirname(filename) != b"":
                    self._log.error(
                        "Only specify a name rather than a path for -n"
                    )
                    return
                for album in lib.albums(args):
                    if opts.associate and (
                        artpath := art.extract_first(
                            self._log,
                            normpath(os.path.join(album.path, filename)),
                            album.items(),
                        )
                    ):
                        album.set_art(artpath)
                        album.store()

        extract_cmd.func = extract_func

        # Clear command.
        clear_cmd = ui.Subcommand(
            "clearart", help="remove images from file metadata"
        )
        clear_cmd.parser.add_option(
            "-y", "--yes", action="store_true", help="skip confirmation"
        )

        def clear_func(
            lib: Library, opts: ClearArtCLIOpts, args: list[str]
        ) -> None:
            items = lib.items(args)
            # Confirm with user.
            if not opts.yes and not _confirm(items, False):
                return
            art.clear(self._log, lib, args)

        clear_cmd.func = clear_func

        return [embed_cmd, extract_cmd, clear_cmd]

    def process_album(self, album: Album) -> None:
        """Automatically embed art after art has been set"""
        if self.config["auto"] and ui.should_write():
            max_width = self.config["maxwidth"].get(int)
            art.embed_album(
                self._log,
                album,
                max_width,
                True,
                self.config["compare_threshold"].get(int),
                self.config["ifempty"].get(bool),
            )
            self.remove_artfile(album)

    def remove_artfile(self, album: Album) -> None:
        """Possibly delete the album art file for an album (if the
        appropriate configuration option is enabled).
        """
        if self.config["remove_art_file"] and album.artpath:
            if os.path.isfile(syspath(album.artpath)):
                self._log.debug("Removing album art file for {}", album)
                os.remove(syspath(album.artpath))
                album.artpath = None
                album.store()

    def import_task_files(
        self, session: ImportSession, task: ImportTask
    ) -> None:
        """Automatically clearart of imported files."""
        for item in task.imported_items():
            self._log.debug("clearart-on-import {.filepath}", item)
            art.clear_item(item, self._log)