ytdl-sub/src/ytdl_sub/plugins/audio_extract.py

import os.path
from typing import Any
from typing import Dict
from typing import Optional

from ytdl_sub.config.plugin import Plugin
from ytdl_sub.config.preset_options import OptionsDictValidator
from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder
from ytdl_sub.entries.entry import Entry
from ytdl_sub.utils.exceptions import FileNotDownloadedException
from ytdl_sub.utils.file_handler import FileMetadata
from ytdl_sub.validators.audo_codec_validator import AUDIO_CODEC_EXTS
from ytdl_sub.validators.audo_codec_validator import AUDIO_CODEC_TYPES_EXTENSION_MAPPING
from ytdl_sub.validators.audo_codec_validator import AudioTypeValidator
from ytdl_sub.validators.validators import FloatValidator


class AudioExtractOptions(OptionsDictValidator):
    """
    Extracts audio from a video file.

    Usage:

    .. code-block:: yaml

       presets:
         my_example_preset:
           audio_extract:
             codec: "mp3"
             quality: 128
    """

    _required_keys = {"codec"}
    _optional_keys = {"quality"}

    @classmethod
    def partial_validate(cls, name: str, value: Any) -> None:
        """
        Partially validate audio extract options
        """
        if isinstance(value, dict):
            value["codec"] = value.get("codec", "mp3")
        _ = cls(name, value)

    def __init__(self, name, value):
        super().__init__(name, value)
        self._codec = self._validate_key(key="codec", validator=AudioTypeValidator).value
        self._quality = self._validate_key_if_present(key="quality", validator=FloatValidator)

    @property
    def codec(self) -> str:
        """
        The codec to output after extracting the audio. Supported codecs are aac, flac, mp3, m4a,
        opus, vorbis, wav, and best to grab the best possible format at runtime.
        """
        return self._codec

    @property
    def quality(self) -> Optional[float]:
        """
        Optional. Specify ffmpeg audio quality. Insert a value between ``0`` (better) and ``9``
        (worse) for variable bitrate, or a specific bitrate like ``128`` for 128k.
        """
        if self._quality is not None:
            return self._quality.value
        return None


class AudioExtractPlugin(Plugin[AudioExtractOptions]):
    plugin_options_type = AudioExtractOptions

    def ytdl_options(self) -> Optional[Dict]:
        """
        Returns
        -------
        YTDL options for extracting audio
        """
        ytdl_options_builder = YTDLOptionsBuilder()

        postprocessor_dict = {
            "key": "FFmpegExtractAudio",
            "when": "post_process",
            "preferredcodec": self.plugin_options.codec,
        }
        if self.plugin_options.quality is not None:
            postprocessor_dict["preferredquality"] = self.plugin_options.quality

        return ytdl_options_builder.add(
            {
                "postprocessors": [postprocessor_dict],
            }
        ).to_dict()

    def modify_entry(self, entry: Entry) -> Optional[Entry]:
        """
        Parameters
        ----------
        entry
            Entry with extracted audio

        Returns
        -------
        Entry with updated 'ext' source variable

        Raises
        ------
        FileNotDownloadedException
            If the audio file is not found
        """
        extracted_audio_file = entry.get_download_file_path()

        # yt-dlp doesn't record which ext is used, so try to find it
        if self.plugin_options.codec == "best":
            new_ext = "m4a"

            for possible_ext in AUDIO_CODEC_EXTS:
                extracted_audio_file = (
                    entry.get_download_file_path().removesuffix(entry.ext) + possible_ext
                )

                if os.path.isfile(extracted_audio_file):
                    new_ext = possible_ext
                    break
        else:
            new_ext = AUDIO_CODEC_TYPES_EXTENSION_MAPPING[self.plugin_options.codec]
            extracted_audio_file = entry.get_download_file_path().removesuffix(entry.ext) + new_ext

        entry.add_kwargs({"ext": new_ext})

        if not self.is_dry_run:
            if not os.path.isfile(extracted_audio_file):
                raise FileNotDownloadedException("Failed to find the extracted audio file")

        return entry

    def post_process_entry(self, entry: Entry) -> Optional[FileMetadata]:
        """
        Warn the user that best cannot infer the format that will be used at run-time.
        """
        if self.plugin_options.codec == "best" and self.is_dry_run:
            return FileMetadata(
                "Caution: extracted audio with 'best' format is not known during dry-run. "
                "Defaulting to m4a."
            )
        return None