From a7f556b58f6bb586342d2cbd292d18f4d82966b0 Mon Sep 17 00:00:00 2001 From: jbannon Date: Wed, 10 Aug 2022 05:40:58 +0000 Subject: [PATCH] new subtitle section. modularize ytdl_options in subscription --- src/ytdl_sub/config/preset.py | 6 + src/ytdl_sub/config/preset_options.py | 85 ++++++++++++-- src/ytdl_sub/downloaders/downloader.py | 13 --- src/ytdl_sub/subscriptions/subscription.py | 45 ++++---- .../subscription_ytdl_options.py | 104 ++++++++++++++++++ src/ytdl_sub/validators/validators.py | 5 + tests/e2e/youtube/test_video.py | 29 +++++ 7 files changed, 237 insertions(+), 50 deletions(-) create mode 100644 src/ytdl_sub/subscriptions/subscription_ytdl_options.py diff --git a/src/ytdl_sub/config/preset.py b/src/ytdl_sub/config/preset.py index b95b589a..4669d34f 100644 --- a/src/ytdl_sub/config/preset.py +++ b/src/ytdl_sub/config/preset.py @@ -14,6 +14,7 @@ from ytdl_sub.config.preset_class_mappings import DownloadStrategyMapping from ytdl_sub.config.preset_class_mappings import PluginMapping from ytdl_sub.config.preset_options import OutputOptions from ytdl_sub.config.preset_options import Overrides +from ytdl_sub.config.preset_options import SubtitleOptions from ytdl_sub.config.preset_options import YTDLOptions from ytdl_sub.downloaders.downloader import Downloader from ytdl_sub.downloaders.downloader import DownloaderValidator @@ -32,6 +33,7 @@ from ytdl_sub.validators.validators import Validator PRESET_KEYS = { "preset", "output_options", + "subtitle_options", "ytdl_options", "overrides", *DownloadStrategyMapping.sources(), @@ -248,6 +250,10 @@ class Preset(StrictDictValidator): validator=OutputOptions, ) + self.subtitle_options = self._validate_key( + key="subtitle_options", validator=SubtitleOptions, default={} + ) + self.ytdl_options = self._validate_key( key="ytdl_options", validator=YTDLOptions, default={} ) diff --git a/src/ytdl_sub/config/preset_options.py b/src/ytdl_sub/config/preset_options.py index a889022d..7e3dee39 100644 --- a/src/ytdl_sub/config/preset_options.py +++ b/src/ytdl_sub/config/preset_options.py @@ -1,3 +1,4 @@ +from typing import List from typing import Optional from yt_dlp.utils import DateRange @@ -9,8 +10,10 @@ from ytdl_sub.validators.string_datetime import StringDatetimeValidator from ytdl_sub.validators.string_formatter_validators import DictFormatterValidator from ytdl_sub.validators.string_formatter_validators import OverridesStringFormatterValidator from ytdl_sub.validators.string_formatter_validators import StringFormatterValidator +from ytdl_sub.validators.string_select_validator import StringSelectValidator from ytdl_sub.validators.validators import BoolValidator from ytdl_sub.validators.validators import LiteralDictValidator +from ytdl_sub.validators.validators import StringListValidator class YTDLOptions(LiteralDictValidator): @@ -132,9 +135,6 @@ class OutputOptions(StrictDictValidator): self._thumbnail_name = self._validate_key_if_present( key="thumbnail_name", validator=StringFormatterValidator ) - self._subtitles_name = self._validate_key_if_present( - key="subtitles_name", validator=StringFormatterValidator - ) self._maintain_download_archive = self._validate_key_if_present( key="maintain_download_archive", validator=BoolValidator, default=False @@ -178,15 +178,6 @@ class OutputOptions(StrictDictValidator): """ return self._thumbnail_name - @property - def subtitles_name(self) -> Optional[StringFormatterValidator]: - """ - Optional. The file name for the media's subtitles if they are present. This can include - directories such as ``"Season {upload_year}/{title}.{subtitle_ext}"``, and will be placed - in the output directory. - """ - return self._subtitles_name - @property def maintain_download_archive(self) -> bool: """ @@ -234,3 +225,73 @@ class OutputOptions(StrictDictValidator): end=self.keep_files_before.datetime_str if self.keep_files_before else None, ) return None + + +class SubtitlesTypeValidator(StringSelectValidator): + _expected_value_type_name = "subtitles type" + _select_values = {"srt", "vtt", "ass", "lrc"} + + +class SubtitleOptions(StrictDictValidator): + _optional_keys = { + "subtitles_name", + "subtitles_type", + "embed_subtitles", + "languages", + "allow_auto_generated_subtitles", + } + + def __init__(self, name, value): + super().__init__(name, value) + self._subtitles_name = self._validate_key_if_present( + key="subtitles_name", validator=StringFormatterValidator + ) + self._subtitles_type = self._validate_key_if_present( + key="subtitles_type", validator=SubtitlesTypeValidator, default="srt" + ) + self._embed_subtitles = self._validate_key_if_present( + key="embed_subtitles", validator=BoolValidator + ) + self._languages = self._validate_key_if_present( + key="languages", validator=StringListValidator, default=["en"] + ).list + self._allow_auto_generated_subtitles = self._validate_key_if_present( + key="allow_auto_generated_subtitles", validator=BoolValidator, default=False + ).value + + @property + def subtitles_name(self) -> Optional[StringFormatterValidator]: + """ + Optional. The file name for the media's subtitles if they are present. This can include + directories such as ``"Season {upload_year}/{title}.{subtitle_ext}"``, and will be placed + in the output directory. + """ + return self._subtitles_name + + @property + def subtitles_type(self) -> Optional[str]: + """ + Optional. The subtitles file format. Defaults to ``srt`` + """ + return self._subtitles_type + + @property + def embed_subtitles(self) -> Optional[bool]: + """ + Optional. Whether to embed the subtitles into the video file. + """ + return self._subtitles_type + + @property + def languages(self) -> Optional[List[str]]: + """ + Optional. Language(s) to download for subtitles. Defaults to ``en`` + """ + return [lang.value for lang in self._languages] + + @property + def allow_auto_generated_subtitles(self) -> Optional[bool]: + """ + Optional. Whether to allow auto generated subtitles. Defaults to False. + """ + return self._allow_auto_generated_subtitles diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index d9ea96b7..d7fe00ba 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -60,11 +60,6 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] _extract_entry_num_retries: int = 5 _extract_entry_retry_wait_sec: int = 3 - @classmethod - def ytdl_option_overrides(cls) -> Dict: - """Global overrides that even overwrite user input""" - return {} - @classmethod def ytdl_option_defaults(cls) -> Dict: """ @@ -78,7 +73,6 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] @classmethod def _configure_ytdl_options( cls, - working_directory: str, ytdl_options: Optional[Dict], ) -> Dict: """Configure the ytdl options for the downloader""" @@ -88,12 +82,6 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] # Overwrite defaults with input ytdl_options = dict(cls.ytdl_option_defaults(), **ytdl_options) - # Overwrite defaults + input with global options - ytdl_options = dict(ytdl_options, **cls.ytdl_option_overrides()) - - # Overwrite the output location with the specified working directory - ytdl_options["outtmpl"] = str(Path(working_directory) / "%(id)s.%(ext)s") - return ytdl_options def __init__( @@ -116,7 +104,6 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] self.download_options = download_options self.ytdl_options = self._configure_ytdl_options( ytdl_options=ytdl_options, - working_directory=self.working_directory, ) @contextmanager diff --git a/src/ytdl_sub/subscriptions/subscription.py b/src/ytdl_sub/subscriptions/subscription.py index acef9971..94420022 100644 --- a/src/ytdl_sub/subscriptions/subscription.py +++ b/src/ytdl_sub/subscriptions/subscription.py @@ -1,5 +1,4 @@ import contextlib -import copy import os import shutil from pathlib import Path @@ -14,12 +13,14 @@ from ytdl_sub.config.config_file import ConfigOptions from ytdl_sub.config.preset import Preset from ytdl_sub.config.preset_options import OutputOptions from ytdl_sub.config.preset_options import Overrides +from ytdl_sub.config.preset_options import SubtitleOptions from ytdl_sub.config.preset_options import YTDLOptions from ytdl_sub.downloaders.downloader import Downloader from ytdl_sub.downloaders.downloader import DownloaderValidator from ytdl_sub.entries.entry import Entry from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.plugin import PluginOptions +from ytdl_sub.subscriptions.subscription_ytdl_options import SubscriptionYTDLOptions from ytdl_sub.utils.file_handler import FileHandlerTransactionLog from ytdl_sub.utils.file_handler import FileMetadata from ytdl_sub.utils.thumbnail import convert_download_thumbnail @@ -109,6 +110,15 @@ class Subscription: """ return self.__preset_options.output_options + @property + def subtitle_options(self) -> SubtitleOptions: + """ + Returns + ------- + The subtitle options defined for this subscription + """ + return self.__preset_options.subtitle_options + @property def overrides(self) -> Overrides: """ @@ -195,7 +205,9 @@ class Subscription: ) # TODO: see if entry even has subtitles - if self.output_options.subtitles_name: + if self.output_options.subtitles_name and ( + entry.kwargs_contains("subtitles") or entry.kwargs_contains("automatic_captions") + ): output_subtitles_name = self.overrides.apply_formatter( formatter=self.output_options.subtitles_name, entry=entry ) @@ -275,29 +287,12 @@ class Subscription: directory. """ self._enhanced_download_archive.reinitialize(dry_run=dry_run) - - # TODO: Move this logic to separate function - # TODO: set id here as well - ytdl_options = copy.deepcopy(self.ytdl_options.dict) - - if self.output_options.thumbnail_name: - ytdl_options["writethumbnail"] = True - if self.output_options.subtitles_name: - ytdl_options["writesubtitles"] = True - ytdl_options["subtitleslangs"] = ["en"] - ytdl_options["postprocessers"] = [ - {"key": "FFmpegSubtitlesConvertorPP", "format": "srt"} - ] - - if dry_run: - ytdl_options["skip_download"] = True - ytdl_options["writethumbnail"] = False - ytdl_options["writesubtitles"] = False - - if self.downloader_class.supports_download_archive and self.maintain_download_archive: - ytdl_options["download_archive"] = str( - Path(self.working_directory) / self._enhanced_download_archive.archive_file_name - ) + ytdl_options = SubscriptionYTDLOptions( + preset=self.__preset_options, + enhanced_download_archive=self._enhanced_download_archive, + working_directory=self.working_directory, + dry_run=dry_run, + ).to_dict() plugins = self._initialize_plugins() with self._subscription_download_context_managers(): diff --git a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py new file mode 100644 index 00000000..c8752149 --- /dev/null +++ b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py @@ -0,0 +1,104 @@ +from pathlib import Path +from typing import Dict + +import mergedeep + +from ytdl_sub.config.preset import Preset +from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive + + +class SubscriptionYTDLOptions: + def __init__( + self, + preset: Preset, + enhanced_download_archive: EnhancedDownloadArchive, + working_directory: str, + dry_run: bool, + ): + self._preset = preset + self._enhanced_download_archive = enhanced_download_archive + self._working_directory = working_directory + self._dry_run = dry_run + + @property + def _global_options(self) -> Dict: + return { + # Download all files in the format of {id}.{ext} + "outtmpl": str(Path(self._working_directory) / "%(id)s.%(ext)s") + } + + @property + def _dry_run_options(self) -> Dict: + return { + "skip_download": True, + "writethumbnail": False, + "writesubtitles": False, + } + + @property + def _output_options(self) -> Dict: + ytdl_options = {} + output_options = self._preset.output_options + + if output_options.thumbnail_name: + ytdl_options["writethumbnail"] = True + if ( + self._preset.downloader.supports_download_archive + and output_options.maintain_download_archive + ): + ytdl_options["download_archive"] = str( + Path(self._working_directory) / self._enhanced_download_archive.archive_file_name + ) + + return ytdl_options + + @property + def _subtitle_options(self) -> Dict: + ytdl_options: Dict = {} + subtitle_options = self._preset.subtitle_options + + write_subtitle_file: bool = subtitle_options.subtitles_name is not None + if write_subtitle_file: + ytdl_options["writesubtitles"] = True + ytdl_options["postprocessors"] = [ + {"key": "FFmpegSubtitlesConvertor", "format": subtitle_options.subtitles_type} + ] + + if subtitle_options.embed_subtitles: + ytdl_options["postprocessors"] = [ + # already_have_subtitle=True means keep the subtitle files. False means delete + {"key": "FFmpegEmbedSubtitle", "already_have_subtitle": write_subtitle_file} + ] + + # If neither subtitles_name or embed_subtitles is set, do not set any other flags + if not ytdl_options: + return {} + + ytdl_options["writeautomaticsub"] = subtitle_options.allow_auto_generated_subtitles + ytdl_options["subtitleslangs"] = subtitle_options.languages + + return ytdl_options + + @property + def _user_ytdl_options(self) -> Dict: + return self._preset.ytdl_options.dict + + def to_dict(self) -> Dict: + ytdl_options = self._global_options + if self._dry_run: + mergedeep.merge( + ytdl_options, + self._dry_run_options, + self._user_ytdl_options, + strategy=mergedeep.Strategy.TYPESAFE_ADDITIVE, + ) + else: + mergedeep.merge( + ytdl_options, + self._output_options, + self._subtitle_options, + self._user_ytdl_options, + strategy=mergedeep.Strategy.TYPESAFE_ADDITIVE, + ) + + return ytdl_options diff --git a/src/ytdl_sub/validators/validators.py b/src/ytdl_sub/validators/validators.py index bc7a29b8..5831ed8c 100644 --- a/src/ytdl_sub/validators/validators.py +++ b/src/ytdl_sub/validators/validators.py @@ -127,6 +127,11 @@ class ListValidator(Validator, ABC, Generic[ValidatorT]): return self._list +class StringListValidator(ListValidator[StringValidator]): + _expected_value_type_name = "string list" + _inner_list_type = StringValidator + + class DictValidator(Validator): """ Validates dictionary-based fields. Errors to them as 'object's since this could be validating diff --git a/tests/e2e/youtube/test_video.py b/tests/e2e/youtube/test_video.py index 4b65df75..2e8b235f 100644 --- a/tests/e2e/youtube/test_video.py +++ b/tests/e2e/youtube/test_video.py @@ -83,6 +83,35 @@ class TestYoutubeVideo: if not dry_run: expected_single_video_download.assert_files_exist(relative_directory=output_directory) + # @pytest.mark.parametrize("dry_run", [True, False]) + # def test_single_video_download_with_subtitles( + # self, + # music_video_config, + # single_video_preset_dict, + # expected_single_video_download, + # output_directory, + # dry_run, + # ): + # single_video_preset_dict["youtube"][ + # "video_url" + # ] = "https://www.youtube.com/watch?v=dQw4w9WgXcQ" + # single_video_preset_dict["subtitles"] = {} + # single_video_preset_dict["subtitles"]["subtitles_name"] = "{music_video_name}.srt" + # single_video_subscription = Subscription.from_dict( + # config=music_video_config, + # preset_name="music_video_single_video_test", + # preset_dict=single_video_preset_dict, + # ) + # + # transaction_log = single_video_subscription.download(dry_run=dry_run) + # assert_transaction_log_matches( + # output_directory=output_directory, + # transaction_log=transaction_log, + # transaction_log_summary_file_name="youtube/test_video.txt", + # ) + # if not dry_run: + # expected_single_video_download.assert_files_exist(relative_directory=output_directory) + @pytest.mark.parametrize("dry_run", [True, False]) def test_single_video_download_from_cli_dl( self,