From 871c051aa8d82fe66064fe622d822f0049da1cd2 Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Sun, 14 Aug 2022 10:10:10 -0700 Subject: [PATCH] [FEATURE] Audio extract plugin (#177) --- docs/config.rst | 8 ++ examples/youtube_extract_and_tag_audio.yaml | 38 ++++++ src/ytdl_sub/config/preset_class_mappings.py | 2 + src/ytdl_sub/downloaders/downloader.py | 2 +- src/ytdl_sub/entries/entry.py | 15 ++- src/ytdl_sub/plugins/audio_extract.py | 108 ++++++++++++++++++ .../subscription_ytdl_options.py | 14 ++- .../validators/audo_codec_validator.py | 22 ++++ src/ytdl_sub/validators/validators.py | 43 +++---- tests/e2e/conftest.py | 10 ++ tests/e2e/plugins/test_audio_extract.py | 91 +++++++++++++++ .../plugins/test_audio_extract_playlist.json | 5 + .../plugins/test_audio_extract_single.json | 3 + .../plugins/test_audio_extract_playlist.txt | 29 +++++ .../plugins/test_audio_extract_single.txt | 11 ++ 15 files changed, 370 insertions(+), 31 deletions(-) create mode 100644 examples/youtube_extract_and_tag_audio.yaml create mode 100644 src/ytdl_sub/plugins/audio_extract.py create mode 100644 src/ytdl_sub/validators/audo_codec_validator.py create mode 100644 tests/e2e/plugins/test_audio_extract.py create mode 100644 tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_playlist.json create mode 100644 tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_single.json create mode 100644 tests/e2e/resources/transaction_log_summaries/plugins/test_audio_extract_playlist.txt create mode 100644 tests/e2e/resources/transaction_log_summaries/plugins/test_audio_extract_single.txt diff --git a/docs/config.rst b/docs/config.rst index dd06cbdd..a646b080 100644 --- a/docs/config.rst +++ b/docs/config.rst @@ -169,6 +169,14 @@ Plugins """"""" Plugins are used to perform any type of post-processing to the already downloaded files. +audio_extract +''''''''''''' +.. autoclass:: ytdl_sub.plugins.audio_extract.AudioExtractOptions() + :members: + :member-order: bysource + +------------------------------------------------------------------------------- + music_tags '''''''''' .. autoclass:: ytdl_sub.plugins.music_tags.MusicTagsOptions() diff --git a/examples/youtube_extract_and_tag_audio.yaml b/examples/youtube_extract_and_tag_audio.yaml new file mode 100644 index 00000000..e738ba59 --- /dev/null +++ b/examples/youtube_extract_and_tag_audio.yaml @@ -0,0 +1,38 @@ +# Extracts audio from YouTube videos, converts it to mp3, and adds audio tags to it. +configuration: + working_directory: '.ytdl-sub-downloads' + +presets: + yt_song: + youtube: + download_strategy: "video" + + output_options: + output_directory: "{music_directory}" + file_name: "{title_sanitized}.{ext}" + + audio_extract: + codec: "mp3" + quality: 128 + + music_tags: + tags: + artist: "{artist}" + albumartist: "{artist}" + title: "{title}" + album: "Singles" + track: "1" + year: "{upload_year}" + genre: "Unset" + + overrides: + music_directory: "/path/to/music" + + yt_song_playlist: + preset: yt_song + youtube: + download_strategy: "playlist" + + music_tags: + tags: + track: "{playlist_index}" \ No newline at end of file diff --git a/src/ytdl_sub/config/preset_class_mappings.py b/src/ytdl_sub/config/preset_class_mappings.py index 3221160a..d3bd1386 100644 --- a/src/ytdl_sub/config/preset_class_mappings.py +++ b/src/ytdl_sub/config/preset_class_mappings.py @@ -9,6 +9,7 @@ from ytdl_sub.downloaders.youtube.merge_playlist import YoutubeMergePlaylistDown from ytdl_sub.downloaders.youtube.playlist import YoutubePlaylistDownloader from ytdl_sub.downloaders.youtube.split_video import YoutubeSplitVideoDownloader from ytdl_sub.downloaders.youtube.video import YoutubeVideoDownloader +from ytdl_sub.plugins.audio_extract import AudioExtractPlugin from ytdl_sub.plugins.music_tags import MusicTagsPlugin from ytdl_sub.plugins.nfo_tags import NfoTagsPlugin from ytdl_sub.plugins.output_directory_nfo_tags import OutputDirectoryNfoTagsPlugin @@ -107,6 +108,7 @@ class PluginMapping: """ _MAPPING: Dict[str, Type[Plugin]] = { + "audio_extract": AudioExtractPlugin, "music_tags": MusicTagsPlugin, "video_tags": VideoTagsPlugin, "nfo_tags": NfoTagsPlugin, diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index e8a872a9..ae808f95 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -132,7 +132,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] def extract_info_with_retry( self, - is_downloaded_fn: Optional[Callable[[], bool]], + is_downloaded_fn: Optional[Callable[[], bool]] = None, ytdl_options_overrides: Optional[Dict] = None, **kwargs, ) -> Dict: diff --git a/src/ytdl_sub/entries/entry.py b/src/ytdl_sub/entries/entry.py index 5970bc26..47236464 100644 --- a/src/ytdl_sub/entries/entry.py +++ b/src/ytdl_sub/entries/entry.py @@ -6,6 +6,7 @@ from typing import final from ytdl_sub.entries.base_entry import BaseEntry from ytdl_sub.entries.variables.entry_variables import EntryVariables +from ytdl_sub.validators.audo_codec_validator import AUDIO_CODEC_EXTS class Entry(EntryVariables, BaseEntry): @@ -65,13 +66,17 @@ class Entry(EntryVariables, BaseEntry): ------- True if the file and thumbnail exist locally. False otherwise. """ - thumbnail_exists = ( - os.path.isfile(self.get_download_thumbnail_path()) - or self.get_ytdlp_download_thumbnail_path() is not None - ) file_exists = os.path.isfile(self.get_download_file_path()) - return file_exists and thumbnail_exists + # HACK: yt-dlp does not record extracted audio extensions anywhere. If the file is not + # found, try it using audio extensions + if not file_exists: + for audio_ext in AUDIO_CODEC_EXTS: + if os.path.isfile(self.get_download_file_path().removesuffix(self.ext) + audio_ext): + file_exists = True + break + + return file_exists @final def to_dict(self) -> Dict[str, str]: diff --git a/src/ytdl_sub/plugins/audio_extract.py b/src/ytdl_sub/plugins/audio_extract.py new file mode 100644 index 00000000..a26d4bfb --- /dev/null +++ b/src/ytdl_sub/plugins/audio_extract.py @@ -0,0 +1,108 @@ +import os.path +from typing import Dict +from typing import Optional + +from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder +from ytdl_sub.entries.entry import Entry +from ytdl_sub.plugins.plugin import Plugin +from ytdl_sub.plugins.plugin import PluginOptions +from ytdl_sub.utils.exceptions import FileNotDownloadedException +from ytdl_sub.validators.audo_codec_validator import AUDIO_CODEC_TYPES_EXTENSION_MAPPING +from ytdl_sub.validators.audo_codec_validator import CodecTypeValidator +from ytdl_sub.validators.validators import FloatValidator + + +class AudioExtractOptions(PluginOptions): + """ + Extracts audio from a video file. + + Usage: + + .. code-block:: yaml + + presets: + my_example_preset: + audio_extract: + codec: "mp3" + quality: 128 + """ + + _optional_keys = {"codec", "quality"} + + def __init__(self, name, value): + super().__init__(name, value) + self._codec = self._validate_key(key="codec", validator=CodecTypeValidator).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. + """ + 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 + """ + new_ext = AUDIO_CODEC_TYPES_EXTENSION_MAPPING[self.plugin_options.codec] + extracted_audio_file = entry.get_download_file_path().removesuffix(entry.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") + + # TODO: create entry function to update kwargs + # pylint: disable=protected-access + entry._kwargs["ext"] = new_ext + # pylint: enable=protected-access + + return entry diff --git a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py index 22b83c02..e4dc95a3 100644 --- a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py +++ b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py @@ -8,6 +8,7 @@ from typing import TypeVar from ytdl_sub.config.preset import Preset from ytdl_sub.downloaders.downloader import Downloader from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder +from ytdl_sub.plugins.audio_extract import AudioExtractPlugin from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.subtitles import SubtitleOptions from ytdl_sub.plugins.subtitles import SubtitlesPlugin @@ -81,6 +82,13 @@ class SubscriptionYTDLOptions: return ytdl_options + @property + def _audio_extract_options(self) -> Dict: + if not (audio_extract_plugin := self._get_plugin(AudioExtractPlugin)): + return {} + + return audio_extract_plugin.ytdl_options() + @property def _subtitle_options(self) -> Dict: if not (subtitle_plugin := self._get_plugin(SubtitlesPlugin)): @@ -90,6 +98,7 @@ class SubscriptionYTDLOptions: # TODO: warn here return {} + # TODO: Use subtitle ytdl_options builder = YTDLOptionsBuilder().add({"writesubtitles": True}) subtitle_options: SubtitleOptions = subtitle_plugin.plugin_options @@ -141,7 +150,10 @@ class SubscriptionYTDLOptions: ) else: ytdl_options_builder.add( - self._output_options, self._subtitle_options, self._user_ytdl_options + self._output_options, + self._subtitle_options, + self._audio_extract_options, + self._user_ytdl_options, ) return ytdl_options_builder diff --git a/src/ytdl_sub/validators/audo_codec_validator.py b/src/ytdl_sub/validators/audo_codec_validator.py new file mode 100644 index 00000000..b986c02f --- /dev/null +++ b/src/ytdl_sub/validators/audo_codec_validator.py @@ -0,0 +1,22 @@ +from typing import Dict +from typing import Set + +from ytdl_sub.validators.string_select_validator import StringSelectValidator + +AUDIO_CODEC_TYPES_EXTENSION_MAPPING: Dict[str, str] = { + "aac": "aac", + "flac": "flac", + "mp3": "mp3", + "m4a": "m4a", + "opus": "opus", + "vorbis": "ogg", + "wav": "wav", +} + +AUDIO_CODEC_TYPES: Set[str] = set(AUDIO_CODEC_TYPES_EXTENSION_MAPPING.keys()) +AUDIO_CODEC_EXTS: Set[str] = set(AUDIO_CODEC_TYPES_EXTENSION_MAPPING.values()) + + +class CodecTypeValidator(StringSelectValidator): + _expected_value_type_name = "codec" + _select_values = AUDIO_CODEC_TYPES diff --git a/src/ytdl_sub/validators/validators.py b/src/ytdl_sub/validators/validators.py index 5831ed8c..694ed6c7 100644 --- a/src/ytdl_sub/validators/validators.py +++ b/src/ytdl_sub/validators/validators.py @@ -11,6 +11,7 @@ from typing import final from ytdl_sub.utils.exceptions import ValidationException +ValueT = TypeVar("ValueT", bound=object) ValidationExceptionT = TypeVar("ValidationExceptionT", bound=ValidationException) ValidatorT = TypeVar("ValidatorT", bound="Validator") @@ -59,40 +60,34 @@ class Validator(ABC): return exception_class(f"{prefix}{error_message}") -class BoolValidator(Validator): +class ValueValidator(Validator, ABC, Generic[ValueT]): """ - Validates boolean fields. + Native type validator that returns the value as-is """ + @property + def value(self) -> ValueT: + """ + Returns + ------- + The value, unmodified + """ + return self._value + + +class BoolValidator(ValueValidator[bool]): _expected_value_type: Type = bool _expected_value_type_name = "boolean" - @property - def value(self) -> bool: - """ - Returns - ------- - Boolean value - """ - return self._value - - -class StringValidator(Validator): - """ - Validates string fields. - """ +class StringValidator(ValueValidator[str]): _expected_value_type = str _expected_value_type_name = "string" - @property - def value(self) -> str: - """ - Returns - ------- - String value - """ - return self._value + +class FloatValidator(ValueValidator[float]): + _expected_value_type = (int, float) + _expected_value_type_name = "float" class ListValidator(Validator, ABC, Generic[ValidatorT]): diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 3b988eba..bc9a78b4 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -38,6 +38,16 @@ def soundcloud_discography_config(): return ConfigFile.from_file_path(config_path="examples/soundcloud_discography_config.yaml") +@pytest.fixture() +def youtube_audio_config_path(): + return "examples/youtube_extract_and_tag_audio.yaml" + + +@pytest.fixture() +def youtube_audio_config(youtube_audio_config_path): + return ConfigFile.from_file_path(config_path=youtube_audio_config_path) + + @pytest.fixture def timestamps_file_path(): timestamps = [ diff --git a/tests/e2e/plugins/test_audio_extract.py b/tests/e2e/plugins/test_audio_extract.py new file mode 100644 index 00000000..198a3df4 --- /dev/null +++ b/tests/e2e/plugins/test_audio_extract.py @@ -0,0 +1,91 @@ +import mergedeep +import pytest +from e2e.expected_download import assert_expected_downloads +from e2e.expected_transaction_log import assert_transaction_log_matches + +from ytdl_sub.subscriptions.subscription import Subscription + + +@pytest.fixture +def single_song_preset_dict(output_directory): + return { + "preset": "yt_song", + "youtube": {"video_url": "https://www.youtube.com/watch?v=2lAe1cqCOXo"}, + "output_options": {"output_directory": output_directory}, + # download the worst format so it is fast + "ytdl_options": { + "format": "worst[ext=mp4]", + }, + } + + +@pytest.fixture +def multiple_songs_preset_dict(output_directory): + return { + "preset": "yt_song_playlist", + "youtube": {"playlist_url": "https://youtube.com/playlist?list=PL5BC0FC26BECA5A35"}, + "output_options": {"output_directory": output_directory}, + "audio_extract": {"codec": "vorbis", "quality": 140}, + # download the worst format so it is fast + "ytdl_options": { + "format": "worst[ext=mp4]", + }, + } + + +class TestAudioExtract: + @pytest.mark.parametrize("dry_run", [True, False]) + def test_audio_extract_single_song( + self, + youtube_audio_config, + single_song_preset_dict, + output_directory, + dry_run, + ): + subscription = Subscription.from_dict( + config=youtube_audio_config, + preset_name="single_song_test", + preset_dict=single_song_preset_dict, + ) + + transaction_log = subscription.download(dry_run=dry_run) + assert_transaction_log_matches( + output_directory=output_directory, + transaction_log=transaction_log, + transaction_log_summary_file_name="plugins/test_audio_extract_single.txt", + regenerate_transaction_log=True, + ) + assert_expected_downloads( + output_directory=output_directory, + dry_run=dry_run, + expected_download_summary_file_name="plugins/test_audio_extract_single.json", + regenerate_expected_download_summary=True, + ) + + @pytest.mark.parametrize("dry_run", [True, False]) + def test_audio_extract_multiple_songs( + self, + youtube_audio_config, + multiple_songs_preset_dict, + output_directory, + dry_run, + ): + subscription = Subscription.from_dict( + config=youtube_audio_config, + preset_name="multiple_songs_test", + preset_dict=multiple_songs_preset_dict, + ) + + transaction_log = subscription.download(dry_run=dry_run) + assert_transaction_log_matches( + output_directory=output_directory, + transaction_log=transaction_log, + transaction_log_summary_file_name="plugins/test_audio_extract_playlist.txt", + regenerate_transaction_log=True, + ) + assert_expected_downloads( + output_directory=output_directory, + dry_run=dry_run, + expected_download_summary_file_name="plugins/test_audio_extract_playlist.json", + regenerate_expected_download_summary=True, + ) diff --git a/tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_playlist.json b/tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_playlist.json new file mode 100644 index 00000000..35c71317 --- /dev/null +++ b/tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_playlist.json @@ -0,0 +1,5 @@ +{ + "Jesse's Minecraft Server [Trailer - Feb.1].ogg": "67b19b495756fb6dc332f49bfa2957d3", + "Jesse's Minecraft Server [Trailer - Feb.27].ogg": "273c96652e198171ee60345051addc7a", + "Jesse's Minecraft Server [Trailer - Mar.21].ogg": "e91db58c153f0724aaed5ce53de2a736" +} \ No newline at end of file diff --git a/tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_single.json b/tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_single.json new file mode 100644 index 00000000..1a2dbb73 --- /dev/null +++ b/tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_single.json @@ -0,0 +1,3 @@ +{ + "YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.mp3": "c4ce74e2ddc4e6a0ab823b5b622e2275" +} \ No newline at end of file diff --git a/tests/e2e/resources/transaction_log_summaries/plugins/test_audio_extract_playlist.txt b/tests/e2e/resources/transaction_log_summaries/plugins/test_audio_extract_playlist.txt new file mode 100644 index 00000000..b02d60ac --- /dev/null +++ b/tests/e2e/resources/transaction_log_summaries/plugins/test_audio_extract_playlist.txt @@ -0,0 +1,29 @@ +Files created in '{output_directory}' +---------------------------------------- +Jesse's Minecraft Server [Trailer - Feb.1].ogg + Music Tags: + album: Singles + albumartist: Project Zombie + artist: Project Zombie + genre: Unset + title: Jesse's Minecraft Server [Trailer - Feb.1] + track: 3 + year: 2011 +Jesse's Minecraft Server [Trailer - Feb.27].ogg + Music Tags: + album: Singles + albumartist: Project Zombie + artist: Project Zombie + genre: Unset + title: Jesse's Minecraft Server [Trailer - Feb.27] + track: 2 + year: 2011 +Jesse's Minecraft Server [Trailer - Mar.21].ogg + Music Tags: + album: Singles + albumartist: Project Zombie + artist: Project Zombie + genre: Unset + title: Jesse's Minecraft Server [Trailer - Mar.21] + track: 1 + year: 2011 \ No newline at end of file diff --git a/tests/e2e/resources/transaction_log_summaries/plugins/test_audio_extract_single.txt b/tests/e2e/resources/transaction_log_summaries/plugins/test_audio_extract_single.txt new file mode 100644 index 00000000..2ee2d4a4 --- /dev/null +++ b/tests/e2e/resources/transaction_log_summaries/plugins/test_audio_extract_single.txt @@ -0,0 +1,11 @@ +Files created in '{output_directory}' +---------------------------------------- +YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.mp3 + Music Tags: + album: Singles + albumartist: YouTube + artist: YouTube + genre: Unset + title: YouTube Rewind 2019: For the Record | #YouTubeRewind + track: 1 + year: 2019 \ No newline at end of file