diff --git a/src/ytdl_sub/config/preset.py b/src/ytdl_sub/config/preset.py index 210d291a..48e4f461 100644 --- a/src/ytdl_sub/config/preset.py +++ b/src/ytdl_sub/config/preset.py @@ -1,6 +1,7 @@ import copy from typing import Any from typing import Dict +from typing import Iterable from typing import List from typing import Optional from typing import Tuple @@ -14,7 +15,6 @@ 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 @@ -33,7 +33,6 @@ from ytdl_sub.validators.validators import Validator PRESET_KEYS = { "preset", "output_options", - "subtitle_options", "ytdl_options", "overrides", *DownloadStrategyMapping.sources(), @@ -41,6 +40,28 @@ PRESET_KEYS = { } +class PresetPlugins: + def __init__(self): + self.plugin_types: List[Type[Plugin]] = [] + self.plugin_options: List[PluginOptions] = [] + + def add(self, plugin_type: Type[Plugin], plugin_options: PluginOptions) -> "PresetPlugins": + """ + Add a pair of plugin type and options to the list + """ + self.plugin_types.append(plugin_type) + self.plugin_options.append(plugin_options) + return self + + def zipped(self) -> Iterable[Tuple[Type[Plugin], PluginOptions]]: + """ + Returns + ------- + Plugin and PluginOptions zipped + """ + return zip(self.plugin_types, self.plugin_options) + + class DownloadStrategyValidator(StrictDictValidator): """ Ensures a download strategy exists for a source. Does not validate any more than that. @@ -79,8 +100,6 @@ class DownloadStrategyValidator(StrictDictValidator): raise self._validation_exception(error_message=value_exc) -# Preset is the meat of ytdl-sub, the linter is wrong here -# pylint: disable=too-many-instance-attributes class Preset(StrictDictValidator): # Have all present keys optional since parent presets could not have all the # required keys. They will get validated in the init after the mergedeep of dicts @@ -140,8 +159,8 @@ class Preset(StrictDictValidator): return downloader, download_options - def __validate_and_get_plugins(self) -> List[Tuple[Type[Plugin], PluginOptions]]: - plugins: List[Tuple[Type[Plugin], PluginOptions]] = [] + def __validate_and_get_plugins(self) -> PresetPlugins: + plugins = PresetPlugins() for key in self._keys: if key not in PluginMapping.plugins(): @@ -153,7 +172,7 @@ class Preset(StrictDictValidator): source_variables=self._source_variables, override_variables=self.overrides.keys ) - plugins.append((plugin, plugin_options)) + plugins.add(plugin_type=plugin, plugin_options=plugin_options) return plugins @@ -171,7 +190,7 @@ class Preset(StrictDictValidator): variable_dict = dict(source_variables, **variable_dict) # For all plugins, add in any extra added source variables - for _, plugin_options in self.plugins: + for plugin_options in self.plugins.plugin_options: added_plugin_variables = { source_var: "dummy_string" for source_var in plugin_options.added_source_variables() } @@ -252,16 +271,12 @@ class Preset(StrictDictValidator): validator=OutputOptions, ) - self.subtitle_options = self._validate_key_if_present( - key="subtitle_options", validator=SubtitleOptions, default={} - ) - self.ytdl_options = self._validate_key( key="ytdl_options", validator=YTDLOptions, default={} ) self.overrides = self._validate_key(key="overrides", validator=Overrides, default={}) - self.plugins = self.__validate_and_get_plugins() + self.plugins: PresetPlugins = self.__validate_and_get_plugins() # After all options are initialized, perform a recursive post-validate that requires # values from multiple validators diff --git a/src/ytdl_sub/config/preset_class_mappings.py b/src/ytdl_sub/config/preset_class_mappings.py index e6d7bf2d..59520211 100644 --- a/src/ytdl_sub/config/preset_class_mappings.py +++ b/src/ytdl_sub/config/preset_class_mappings.py @@ -14,6 +14,7 @@ from ytdl_sub.plugins.nfo_tags import NfoTagsPlugin from ytdl_sub.plugins.output_directory_nfo_tags import OutputDirectoryNfoTagsPlugin from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.regex import RegexPlugin +from ytdl_sub.plugins.subtitles import SubtitlesPlugin class DownloadStrategyMapping: @@ -109,6 +110,7 @@ class PluginMapping: "nfo_tags": NfoTagsPlugin, "output_directory_nfo_tags": OutputDirectoryNfoTagsPlugin, "regex": RegexPlugin, + "subtitles": SubtitlesPlugin, } @classmethod diff --git a/src/ytdl_sub/config/preset_options.py b/src/ytdl_sub/config/preset_options.py index 03e60f47..c2df0ff9 100644 --- a/src/ytdl_sub/config/preset_options.py +++ b/src/ytdl_sub/config/preset_options.py @@ -1,4 +1,3 @@ -from typing import List from typing import Optional from yt_dlp.utils import DateRange @@ -10,10 +9,8 @@ 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): @@ -225,93 +222,3 @@ 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): - """ - Defines how to download and store subtitles. - - Usage: - - .. code-block:: yaml - - presets: - my_example_preset: - subtitle_options: - # required - output_directory: "/path/to/videos_or_music" - file_name: "{title_sanitized}.{ext}" - # optional - thumbnail_name: "{title_sanitized}.{thumbnail_ext}" - maintain_download_archive: True - keep_files_before: now - keep_files_after: 19000101 - """ - - _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/entries/variables/entry_variables.py b/src/ytdl_sub/entries/variables/entry_variables.py index 1c2ff226..fccf892f 100644 --- a/src/ytdl_sub/entries/variables/entry_variables.py +++ b/src/ytdl_sub/entries/variables/entry_variables.py @@ -118,8 +118,7 @@ class EntryVariables(SourceVariables): Returns ------- str - The download entry's thumbnail extension. Will always return 'str'. Until there is a - need to support other subtitle types, we always use srt. + The download entry's subtitles extension. """ return "srt" diff --git a/src/ytdl_sub/plugins/plugin.py b/src/ytdl_sub/plugins/plugin.py index 6283f758..654957e0 100644 --- a/src/ytdl_sub/plugins/plugin.py +++ b/src/ytdl_sub/plugins/plugin.py @@ -1,4 +1,5 @@ from abc import ABC +from typing import Dict from typing import Generic from typing import List from typing import Optional @@ -73,6 +74,13 @@ class Plugin(DownloadArchiver, Generic[PluginOptionsT], ABC): # TODO pass yaml snake case name in the class somewhere, and use it for the logger self._logger = Logger.get(self.__class__.__name__) + def ytdl_options(self) -> Optional[Dict]: + """ + Returns + ------- + ytdl options to enable/disable when downloading entries for this specific plugin + """ + # pylint: disable=no-self-use def modify_entry(self, entry: Entry) -> Optional[Entry]: """ diff --git a/src/ytdl_sub/plugins/subtitles.py b/src/ytdl_sub/plugins/subtitles.py new file mode 100644 index 00000000..26d00589 --- /dev/null +++ b/src/ytdl_sub/plugins/subtitles.py @@ -0,0 +1,170 @@ +from typing import Dict +from typing import List +from typing import Optional +from typing import Set + +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.validators.strict_dict_validator import StrictDictValidator +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 StringListValidator + +SUBTITLE_EXTENSIONS: Set[str] = {"srt", "vtt", "ass", "lrc"} + + +class SubtitlesTypeValidator(StringSelectValidator): + _expected_value_type_name = "subtitles type" + _select_values = SUBTITLE_EXTENSIONS + + +class SubtitleOptions(StrictDictValidator): + """ + Defines how to download and store subtitles. + + Usage: + + .. code-block:: yaml + + presets: + my_example_preset: + subtitle_options: + # required + output_directory: "/path/to/videos_or_music" + file_name: "{title_sanitized}.{ext}" + # optional + thumbnail_name: "{title_sanitized}.{thumbnail_ext}" + maintain_download_archive: True + keep_files_before: now + keep_files_after: 19000101 + """ + + _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 + + +class SubtitlesPlugin(Plugin[SubtitleOptions]): + plugin_options_type = SubtitleOptions + + def ytdl_options(self) -> Optional[Dict]: + ytdl_options_builder = YTDLOptionsBuilder() + + write_subtitle_file: bool = self.plugin_options.subtitles_name is not None + if write_subtitle_file: + ytdl_options_builder.add( + { + "writesubtitles": True, + "postprocessors": { + "key": "FFmpegSubtitlesConvertor", + "format": self.plugin_options.subtitles_type, + }, + } + ) + + if self.plugin_options.embed_subtitles: + ytdl_options_builder.add( + { + "postprocessors": [ + # already_have_subtitle=True means keep the subtitle files + {"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_builder.to_dict(): + return {} + + return ytdl_options_builder.add( + { + "writeautomaticsub": self.plugin_options.allow_auto_generated_subtitles, + "subtitleslangs": self.plugin_options.languages, + } + ).to_dict() + + def post_process_entry(self, entry: Entry) -> None: + """ + Creates an entry's NFO file using values defined in the metadata options + + Parameters + ---------- + entry: + Entry to create an NFO file for + """ + + # def get_ytdlp_download_subtitle_paths(self) -> List[str]: + # possible_subtitle_exts = SUBTITLE_EXTENSIONS + # subtitle_paths: List[str] = [] + # + # for ext in possible_subtitle_exts: + # for path in Path(self.working_directory()).rglob("*"): + # if ( + # path.is_file() + # and path.name.startswith(self.uid) + # and path.name.endswith(f".{ext}") + # ): + # subtitle_paths.append(str(path)) + # + # return subtitle_paths diff --git a/src/ytdl_sub/subscriptions/subscription.py b/src/ytdl_sub/subscriptions/subscription.py index 5941916f..bcd1261c 100644 --- a/src/ytdl_sub/subscriptions/subscription.py +++ b/src/ytdl_sub/subscriptions/subscription.py @@ -5,21 +5,20 @@ from pathlib import Path from typing import Dict from typing import List from typing import Optional -from typing import Tuple from typing import Type from ytdl_sub.config.config_file import ConfigFile from ytdl_sub.config.config_file import ConfigOptions from ytdl_sub.config.preset import Preset +from ytdl_sub.config.preset import PresetPlugins 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.plugins.subtitles import SubtitleOptions 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 @@ -84,7 +83,7 @@ class Subscription: return self.__preset_options.downloader_options @property - def plugins(self) -> List[Tuple[Type[Plugin], PluginOptions]]: + def plugins(self) -> PresetPlugins: """ Returns ------- @@ -204,18 +203,17 @@ class Subscription: entry=entry, ) - # if self.downloader_class.supports_subtitles and self.subtitle_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 - # ) - # - # self._enhanced_download_archive.save_file_to_output_directory( - # file_name=entry.get_download_subtitles_name(), - # output_file_name=output_subtitles_name, - # entry=entry, - # ) + if self.downloader_class.supports_subtitles and self.subtitle_options.subtitles_name: + + output_subtitles_name = self.overrides.apply_formatter( + formatter=self.output_options.subtitles_name, entry=entry + ) + + self._enhanced_download_archive.save_file_to_output_directory( + file_name=entry.get_download_subtitles_name(), + output_file_name=output_subtitles_name, + entry=entry, + ) @contextlib.contextmanager def _prepare_working_directory(self): @@ -264,7 +262,7 @@ class Subscription: List of plugins defined in the subscription, initialized and ready to use. """ plugins: List[Plugin] = [] - for plugin_type, plugin_options in self.plugins: + for plugin_type, plugin_options in self.plugins.zipped(): plugin = plugin_type( plugin_options=plugin_options, overrides=self.overrides, @@ -286,14 +284,16 @@ class Subscription: directory. """ self._enhanced_download_archive.reinitialize(dry_run=dry_run) + plugins = self._initialize_plugins() + ytdl_options_builder = SubscriptionYTDLOptions( preset=self.__preset_options, + plugins=plugins, enhanced_download_archive=self._enhanced_download_archive, working_directory=self.working_directory, dry_run=dry_run, ).builder() - plugins = self._initialize_plugins() with self._subscription_download_context_managers(): downloader = self.downloader_class( download_options=self.downloader_options, diff --git a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py index 46360d93..a4b9436d 100644 --- a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py +++ b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py @@ -1,26 +1,42 @@ from pathlib import Path from typing import Dict +from typing import List +from typing import Optional from typing import Type +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.plugin import Plugin +from ytdl_sub.plugins.subtitles import SubtitleOptions +from ytdl_sub.plugins.subtitles import SubtitlesPlugin from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive +PluginT = TypeVar("PluginT", bound=Plugin) + class SubscriptionYTDLOptions: def __init__( self, preset: Preset, + plugins: List[Plugin], enhanced_download_archive: EnhancedDownloadArchive, working_directory: str, dry_run: bool, ): self._preset = preset + self._plugins = plugins self._enhanced_download_archive = enhanced_download_archive self._working_directory = working_directory self._dry_run = dry_run + def _get_plugin(self, plugin_type: Type[PluginT]) -> Optional[PluginT]: + for plugin in self._plugins: + if isinstance(plugin, plugin_type): + return plugin + return None + @property def _downloader(self) -> Type[Downloader]: return self._preset.downloader @@ -67,11 +83,15 @@ class SubscriptionYTDLOptions: @property def _subtitle_options(self) -> Dict: + if not (subtitle_plugin := self._get_plugin(SubtitlesPlugin)): + return {} + if not self._downloader.supports_subtitles: + # TODO: warn here return {} ytdl_options: Dict = {} - subtitle_options = self._preset.subtitle_options + subtitle_options: SubtitleOptions = subtitle_plugin.plugin_options write_subtitle_file: bool = subtitle_options.subtitles_name is not None if write_subtitle_file: