diff --git a/src/ytdl_sub/plugins/audio_extract.py b/src/ytdl_sub/plugins/audio_extract.py index 7418a146..750429de 100644 --- a/src/ytdl_sub/plugins/audio_extract.py +++ b/src/ytdl_sub/plugins/audio_extract.py @@ -8,6 +8,8 @@ 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 @@ -49,7 +51,7 @@ class AudioExtractOptions(OptionsDictValidator): def codec(self) -> str: """ The codec to output after extracting the audio. Supported codecs are aac, flac, mp3, m4a, - opus, vorbis, wav. + opus, vorbis, wav, and best to grab the best possible format at runtime. """ return self._codec @@ -105,12 +107,39 @@ class AudioExtractPlugin(Plugin[AudioExtractOptions]): 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 + 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") - entry.add_kwargs({"ext": new_ext}) - 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 diff --git a/src/ytdl_sub/validators/audo_codec_validator.py b/src/ytdl_sub/validators/audo_codec_validator.py index 7410ac86..959d697a 100644 --- a/src/ytdl_sub/validators/audo_codec_validator.py +++ b/src/ytdl_sub/validators/audo_codec_validator.py @@ -21,7 +21,7 @@ VIDEO_CODEC_EXTS: Set[str] = {"avi", "flv", "mkv", "mov", "mp4", "webm"} class AudioTypeValidator(StringSelectValidator): _expected_value_type_name = "codec" - _select_values = AUDIO_CODEC_TYPES + _select_values = AUDIO_CODEC_TYPES.union({"best"}) # support 'best' in the audio extract plugin class FileTypeValidator(StringSelectValidator): diff --git a/tests/e2e/plugins/test_audio_extract.py b/tests/e2e/plugins/test_audio_extract.py index ef906929..4f6c3e56 100644 --- a/tests/e2e/plugins/test_audio_extract.py +++ b/tests/e2e/plugins/test_audio_extract.py @@ -43,6 +43,11 @@ def single_song_preset_dict(output_directory): } +@pytest.fixture +def single_song_best_format_preset_dict(single_song_preset_dict): + return dict(single_song_preset_dict, **{"audio_extract": {"codec": "best"}}) + + @pytest.fixture def multiple_songs_preset_dict(output_directory): return { @@ -88,7 +93,7 @@ class TestAudioExtract: ) @pytest.mark.parametrize("dry_run", [True, False]) - def test_audio_extract_single_song_new_format( + def test_audio_extract_single_song( self, music_audio_config, single_song_preset_dict, @@ -113,6 +118,32 @@ class TestAudioExtract: expected_download_summary_file_name="plugins/test_audio_extract_single.json", ) + @pytest.mark.parametrize("dry_run", [True, False]) + def test_audio_extract_single_song_best_format( + self, + music_audio_config, + single_song_best_format_preset_dict, + output_directory, + dry_run, + ): + subscription = Subscription.from_dict( + config=music_audio_config, + preset_name="single_song_best_test", + preset_dict=single_song_best_format_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=f"plugins/test_audio_extract_single_best{'_dry_run' if dry_run else ''}.txt", + ) + assert_expected_downloads( + output_directory=output_directory, + dry_run=dry_run, + expected_download_summary_file_name=f"plugins/test_audio_extract_single_best{'_dry_run' if dry_run else ''}.json", + ) + @pytest.mark.parametrize("dry_run", [True, False]) def test_audio_extract_multiple_songs( self, diff --git a/tests/resources/expected_downloads_summaries/plugins/test_audio_extract_single_best.json b/tests/resources/expected_downloads_summaries/plugins/test_audio_extract_single_best.json new file mode 100644 index 00000000..5632536f --- /dev/null +++ b/tests/resources/expected_downloads_summaries/plugins/test_audio_extract_single_best.json @@ -0,0 +1,5 @@ +{ + ".ytdl-sub-single_song_best_test-download-archive.json": "f179ccfe0a9b3a76ae62122a3ccb58fd", + "YouTube/[2019] YouTube Rewind 2019： For the Record ｜ #YouTubeRewind/01 - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.m4a": "4b03133e9422a9ce196d271beb15ceef", + "YouTube/[2019] YouTube Rewind 2019： For the Record ｜ #YouTubeRewind/folder.jpg": "50ee47c80f679029f5d3503bb91b045a" +} \ No newline at end of file diff --git a/tests/resources/transaction_log_summaries/plugins/test_audio_extract_single_best.txt b/tests/resources/transaction_log_summaries/plugins/test_audio_extract_single_best.txt new file mode 100644 index 00000000..9ae37e69 --- /dev/null +++ b/tests/resources/transaction_log_summaries/plugins/test_audio_extract_single_best.txt @@ -0,0 +1,22 @@ +Files created: +---------------------------------------- +{output_directory} + .ytdl-sub-single_song_best_test-download-archive.json +{output_directory}/YouTube/[2019] YouTube Rewind 2019： For the Record ｜ #YouTubeRewind + 01 - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.m4a + Music Tags: + album: YouTube Rewind 2019: For the Record | #YouTubeRewind + albumartist: YouTube + albumartists: YouTube + artist: YouTube + artists: YouTube + genre: Unset + genres: + - multi_tag_1 + - multi_tag_2 + title: YouTube Rewind 2019: For the Record | #YouTubeRewind + track: 1 + tracktotal: 1 + year: 2019 + Embedded thumbnail + folder.jpg \ No newline at end of file diff --git a/tests/resources/transaction_log_summaries/plugins/test_audio_extract_single_best_dry_run.txt b/tests/resources/transaction_log_summaries/plugins/test_audio_extract_single_best_dry_run.txt new file mode 100644 index 00000000..dcc32ffa --- /dev/null +++ b/tests/resources/transaction_log_summaries/plugins/test_audio_extract_single_best_dry_run.txt @@ -0,0 +1,23 @@ +Files created: +---------------------------------------- +{output_directory} + .ytdl-sub-single_song_best_test-download-archive.json +{output_directory}/YouTube/[2019] YouTube Rewind 2019： For the Record ｜ #YouTubeRewind + 01 - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.m4a + Caution: extracted audio with 'best' format is not known during dry-run. Defaulting to m4a. + Music Tags: + album: YouTube Rewind 2019: For the Record | #YouTubeRewind + albumartist: YouTube + albumartists: YouTube + artist: YouTube + artists: YouTube + genre: Unset + genres: + - multi_tag_1 + - multi_tag_2 + title: YouTube Rewind 2019: For the Record | #YouTubeRewind + track: 1 + tracktotal: 1 + year: 2019 + Embedded thumbnail + folder.jpg \ No newline at end of file