From 1f09398ad563b86e1d5b9c794d4f26dde7a180ff Mon Sep 17 00:00:00 2001 From: Jesse Bannon Date: Fri, 12 Aug 2022 09:58:08 -0700 Subject: [PATCH] [FEATURE] Subtitles support via Subtitles plugin (#169) --- docs/config.rst | 8 + src/ytdl_sub/config/preset.py | 33 ++- src/ytdl_sub/config/preset_class_mappings.py | 2 + src/ytdl_sub/config/preset_options.py | 24 +- src/ytdl_sub/downloaders/downloader.py | 72 +++--- .../soundcloud/albums_and_singles.py | 1 + src/ytdl_sub/downloaders/youtube/channel.py | 5 +- .../downloaders/youtube/merge_playlist.py | 1 + .../downloaders/youtube/split_video.py | 3 + .../downloaders/ytdl_options_builder.py | 59 +++++ src/ytdl_sub/plugins/plugin.py | 8 + src/ytdl_sub/plugins/subtitles.py | 209 ++++++++++++++++++ src/ytdl_sub/subscriptions/subscription.py | 34 ++- .../subscription_ytdl_options.py | 147 ++++++++++++ .../validators/strict_dict_validator.py | 7 - src/ytdl_sub/validators/validators.py | 5 + tests/e2e/plugins/test_subtitles.py | 87 ++++++++ .../plugins/test_subtitles_embedded.json | 5 + .../test_subtitles_embedded_and_file.json | 7 + .../plugins/test_subtitles_embedded.txt | 12 + .../test_subtitles_embedded_and_file.txt | 14 ++ 21 files changed, 662 insertions(+), 81 deletions(-) create mode 100644 src/ytdl_sub/downloaders/ytdl_options_builder.py create mode 100644 src/ytdl_sub/plugins/subtitles.py create mode 100644 src/ytdl_sub/subscriptions/subscription_ytdl_options.py create mode 100644 tests/e2e/plugins/test_subtitles.py create mode 100644 tests/e2e/resources/expected_downloads_summaries/plugins/test_subtitles_embedded.json create mode 100644 tests/e2e/resources/expected_downloads_summaries/plugins/test_subtitles_embedded_and_file.json create mode 100644 tests/e2e/resources/transaction_log_summaries/plugins/test_subtitles_embedded.txt create mode 100644 tests/e2e/resources/transaction_log_summaries/plugins/test_subtitles_embedded_and_file.txt diff --git a/docs/config.rst b/docs/config.rst index 45d2565d..55a2fec3 100644 --- a/docs/config.rst +++ b/docs/config.rst @@ -203,6 +203,14 @@ regex ------------------------------------------------------------------------------- +subtitles +''''''''' +.. autoclass:: ytdl_sub.plugins.subtitles.SubtitleOptions() + :members: subtitles_name, subtitles_type, embed_subtitles, languages, allow_auto_generated_subtitles + :member-order: bysource + +------------------------------------------------------------------------------- + .. _subscription_yaml: subscription.yaml diff --git a/src/ytdl_sub/config/preset.py b/src/ytdl_sub/config/preset.py index b95b589a..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 @@ -39,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. @@ -136,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(): @@ -149,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 @@ -167,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() } @@ -253,7 +276,7 @@ class Preset(StrictDictValidator): ) 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 a7f15ed5..b3bcc53e 100644 --- a/src/ytdl_sub/config/preset_options.py +++ b/src/ytdl_sub/config/preset_options.py @@ -1,3 +1,4 @@ +from typing import Dict from typing import Optional from yt_dlp.utils import DateRange @@ -74,15 +75,31 @@ class Overrides(DictFormatterValidator): ) def apply_formatter( - self, formatter: StringFormatterValidator, entry: Optional[Entry] = None + self, + formatter: StringFormatterValidator, + entry: Optional[Entry] = None, + function_overrides: Dict[str, str] = None, ) -> str: """ - Returns the format_string after .format has been called on it using entry (if provided) and - override values + Parameters + ---------- + formatter + Formatter to apply + entry + Optional. Entry to add source variables to the formatter + function_overrides + Optional. Explicit values to override the overrides themselves and source variables + + Returns + ------- + The format_string after .format has been called """ variable_dict = self.dict_with_format_strings if entry: variable_dict = dict(entry.to_dict(), **variable_dict) + if function_overrides: + variable_dict = dict(variable_dict, **function_overrides) + return formatter.apply_formatter(variable_dict) @@ -110,6 +127,7 @@ class OutputOptions(StrictDictValidator): _required_keys = {"output_directory", "file_name"} _optional_keys = { "thumbnail_name", + "subtitles_name", "maintain_download_archive", "keep_files_before", "keep_files_after", diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 5f19999a..22e83dc7 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -1,6 +1,5 @@ import abc import contextlib -import copy import json import os import time @@ -22,6 +21,7 @@ from yt_dlp.utils import ExistingVideoReached from yt_dlp.utils import RejectedVideoReached from ytdl_sub.config.preset_options import Overrides +from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder from ytdl_sub.entries.base_entry import BaseEntry from ytdl_sub.entries.entry import Entry from ytdl_sub.thread.log_entries_downloaded_listener import LogEntriesDownloadedListener @@ -56,15 +56,11 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] downloader_entry_type: Type[Entry] = Entry supports_download_archive: bool = True + supports_subtitles: bool = True _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: """ @@ -75,32 +71,11 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] """ return {"ignoreerrors": True} - @classmethod - def _configure_ytdl_options( - cls, - working_directory: str, - ytdl_options: Optional[Dict], - ) -> Dict: - """Configure the ytdl options for the downloader""" - if ytdl_options is None: - ytdl_options = {} - - # 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__( self, download_options: DownloaderOptionsT, enhanced_download_archive: EnhancedDownloadArchive, - ytdl_options: Optional[Dict] = None, + ytdl_options_builder: YTDLOptionsBuilder, ): """ Parameters @@ -109,14 +84,14 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] Options validator for this downloader enhanced_download_archive Download archive - ytdl_options - YTDL options validator + ytdl_options_builder + YTDL options builder """ DownloadArchiver.__init__(self=self, enhanced_download_archive=enhanced_download_archive) self.download_options = download_options - self.ytdl_options = self._configure_ytdl_options( - ytdl_options=ytdl_options, - working_directory=self.working_directory, + + self._ytdl_options_builder = ytdl_options_builder.clone().add( + self.ytdl_option_defaults(), before=True ) @contextmanager @@ -124,9 +99,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] """ Context manager to interact with yt_dlp. """ - ytdl_options = self.ytdl_options - if ytdl_options_overrides is not None: - ytdl_options = dict(ytdl_options, **ytdl_options_overrides) + ytdl_options = self._ytdl_options_builder.clone().add(ytdl_options_overrides).to_dict() download_logger.debug("ytdl_options: %s", str(ytdl_options)) with Logger.handle_external_logs(name="yt-dlp"): @@ -140,7 +113,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] ------- True if dry-run is enabled. False otherwise. """ - return self.ytdl_options.get("skip_download", False) + return self._ytdl_options_builder.to_dict().get("skip_download", False) def extract_info(self, ytdl_options_overrides: Optional[Dict] = None, **kwargs) -> Dict: """ @@ -186,7 +159,6 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] """ num_tries = 0 entry_files_exist = False - ytdl_options_overrides = copy.deepcopy(ytdl_options_overrides) while not entry_files_exist and num_tries < self._extract_entry_num_retries: entry_dict = self.extract_info(ytdl_options_overrides=ytdl_options_overrides, **kwargs) @@ -198,7 +170,11 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] # Remove the download archive so it can retry without thinking its already downloaded, # even though it is not - ytdl_options_overrides["download_archive"] = None + ytdl_options_overrides = ( + YTDLOptionsBuilder() + .add(ytdl_options_overrides, {"download_archive": None}) + .to_dict() + ) if num_tries < self._extract_entry_retry_wait_sec: download_logger.debug( @@ -312,19 +288,25 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] **kwargs arguments passed directory to YoutubeDL extract_info """ + ytdl_options_builder = self._ytdl_options_builder.clone() if ytdl_options_overrides is None: ytdl_options_overrides = {} - extract_info_ytdl_options = {"writeinfojson": True} + ytdl_options_builder.add({"writeinfojson": True}, ytdl_options_overrides) if only_info_json: - extract_info_ytdl_options["skip_download"] = True - extract_info_ytdl_options["writethumbnail"] = False - - ytdl_options_overrides = dict(ytdl_options_overrides, **extract_info_ytdl_options) + ytdl_options_builder.add( + { + "skip_download": True, + "writethumbnail": False, + "writesubtitles": False, + } + ) try: with self._listen_and_log_downloaded_info_json(log_prefix=log_prefix_on_info_json_dl): - _ = self.extract_info(ytdl_options_overrides=ytdl_options_overrides, **kwargs) + _ = self.extract_info( + ytdl_options_overrides=ytdl_options_builder.to_dict(), **kwargs + ) except RejectedVideoReached: download_logger.debug("RejectedVideoReached, stopping additional downloads") except ExistingVideoReached: diff --git a/src/ytdl_sub/downloaders/soundcloud/albums_and_singles.py b/src/ytdl_sub/downloaders/soundcloud/albums_and_singles.py index e5e9f760..a2ad1720 100644 --- a/src/ytdl_sub/downloaders/soundcloud/albums_and_singles.py +++ b/src/ytdl_sub/downloaders/soundcloud/albums_and_singles.py @@ -51,6 +51,7 @@ class SoundcloudAlbumsAndSinglesDownloader( SoundcloudDownloader[SoundcloudAlbumsAndSinglesDownloadOptions] ): downloader_options_type = SoundcloudAlbumsAndSinglesDownloadOptions + supports_subtitles = False @classmethod def ytdl_option_defaults(cls) -> Dict: diff --git a/src/ytdl_sub/downloaders/youtube/channel.py b/src/ytdl_sub/downloaders/youtube/channel.py index f132393b..e84056dd 100644 --- a/src/ytdl_sub/downloaders/youtube/channel.py +++ b/src/ytdl_sub/downloaders/youtube/channel.py @@ -9,6 +9,7 @@ from ytdl_sub.downloaders.downloader import DownloaderOptionsT from ytdl_sub.downloaders.downloader import download_logger from ytdl_sub.downloaders.youtube.abc import YoutubeDownloader from ytdl_sub.downloaders.youtube.abc import YoutubeDownloaderOptions +from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder from ytdl_sub.entries.youtube import YoutubeChannel from ytdl_sub.entries.youtube import YoutubeVideo from ytdl_sub.utils.thumbnail import convert_url_thumbnail @@ -115,12 +116,12 @@ class YoutubeChannelDownloader(YoutubeDownloader[YoutubeChannelDownloaderOptions self, download_options: DownloaderOptionsT, enhanced_download_archive: EnhancedDownloadArchive, - ytdl_options: Optional[Dict] = None, + ytdl_options_builder: YTDLOptionsBuilder, ): super().__init__( download_options=download_options, enhanced_download_archive=enhanced_download_archive, - ytdl_options=ytdl_options, + ytdl_options_builder=ytdl_options_builder, ) self.channel: Optional[YoutubeChannel] = None diff --git a/src/ytdl_sub/downloaders/youtube/merge_playlist.py b/src/ytdl_sub/downloaders/youtube/merge_playlist.py index e585a1fc..e71adf04 100644 --- a/src/ytdl_sub/downloaders/youtube/merge_playlist.py +++ b/src/ytdl_sub/downloaders/youtube/merge_playlist.py @@ -64,6 +64,7 @@ class YoutubeMergePlaylistDownloader( downloader_options_type = YoutubeMergePlaylistDownloaderOptions downloader_entry_type = YoutubeVideo supports_download_archive = False + supports_subtitles = False @classmethod def ytdl_option_defaults(cls) -> Dict: diff --git a/src/ytdl_sub/downloaders/youtube/split_video.py b/src/ytdl_sub/downloaders/youtube/split_video.py index c113af63..a7059c74 100644 --- a/src/ytdl_sub/downloaders/youtube/split_video.py +++ b/src/ytdl_sub/downloaders/youtube/split_video.py @@ -107,6 +107,9 @@ class YoutubeSplitVideoDownloader( downloader_options_type = YoutubeSplitVideoDownloaderOptions downloader_entry_type = YoutubePlaylistVideo + supports_download_archive = False + supports_subtitles = False + @classmethod def ytdl_option_defaults(cls) -> Dict: """ diff --git a/src/ytdl_sub/downloaders/ytdl_options_builder.py b/src/ytdl_sub/downloaders/ytdl_options_builder.py new file mode 100644 index 00000000..b159da42 --- /dev/null +++ b/src/ytdl_sub/downloaders/ytdl_options_builder.py @@ -0,0 +1,59 @@ +import copy +from typing import Dict +from typing import Optional + +import mergedeep + + +class YTDLOptionsBuilder: + """ + A class to track any modifications made to ytdl options + """ + + def __init__(self): + self._ytdl_options: Dict = {} + + def add( + self, + *ytdl_option_dicts: Optional[Dict], + before: bool = False, + strategy: mergedeep.Strategy = mergedeep.Strategy.TYPESAFE_ADDITIVE + ) -> "YTDLOptionsBuilder": + """ + Parameters + ---------- + *ytdl_option_dicts + One or many ytdl_option dicts. Can also contain None's for convenience + before + Optional. Whether to add these dicts before or after the original + strategy + Optional. mergedeep strategy. Defaults to TYPESAFE_ADDITIVE + + Returns + ------- + instance with the added ytdl_option dict(s) + """ + non_empty = [ytdl_options for ytdl_options in ytdl_option_dicts if ytdl_options is not None] + + if before: + non_empty.append(self.to_dict()) + self._ytdl_options = {} + + mergedeep.merge(self._ytdl_options, *non_empty, strategy=strategy) + return self + + def clone(self) -> "YTDLOptionsBuilder": + """ + Returns + ------- + Deep-copied instance + """ + return copy.deepcopy(self) + + def to_dict(self) -> Dict: + """ + Returns + ------- + Deep-copied dict of the current builder state + """ + return copy.deepcopy(self._ytdl_options) 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..662a1b39 --- /dev/null +++ b/src/ytdl_sub/plugins/subtitles.py @@ -0,0 +1,209 @@ +from pathlib import Path +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.plugins.plugin import PluginOptions +from ytdl_sub.utils.file_handler import FileMetadata +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"} + + +def _is_entry_subtitle_file(path: Path, entry: Entry) -> bool: + if path.is_file() and path.name.startswith(entry.uid): + for ext in SUBTITLE_EXTENSIONS: + if path.name.endswith(f".{ext}"): + return True + return False + + +class SubtitlesTypeValidator(StringSelectValidator): + _expected_value_type_name = "subtitles type" + _select_values = SUBTITLE_EXTENSIONS + + +class SubtitleOptions(PluginOptions): + """ + Defines how to download and store subtitles. Using this plugin creates two new variables: + ``lang`` and ``subtitles_ext``. ``lang`` is dynamic since you can download multiple subtitles. + It will set the respective language to the correct subtitle file. + + Usage: + + .. code-block:: yaml + + presets: + my_example_preset: + subtitle_options: + subtitles_name: "{title_sanitized}.{lang}.{subtitles_ext}" + subtitles_type: "srt" + embed_subtitles: False + languages: "en" # supports list of multiple languages + allow_auto_generated_subtitles: False + """ + + _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" + ).value + self._embed_subtitles = self._validate_key_if_present( + key="embed_subtitles", validator=BoolValidator + ).value + 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_sanitized}.{lang}.{subtitles_ext}"``, and + will be placed in the output directory. ``lang`` is dynamic since you can download multiple + subtitles. It will set the respective language to the correct subtitle file. + """ + return self._subtitles_name + + @property + def subtitles_type(self) -> Optional[str]: + """ + Optional. One of the subtitle file types "srt", "vtt", "ass", "lrc". Defaults to "srt" + """ + return self._subtitles_type + + @property + def embed_subtitles(self) -> Optional[bool]: + """ + Optional. Whether to embed the subtitles into the video file. Defaults to False. + NOTE: webm files can only embed "vtt" subtitle types. + """ + return self._embed_subtitles + + @property + def languages(self) -> Optional[List[str]]: + """ + Optional. Language code(s) to download for subtitles. Supports a single or list of multiple + language codes. 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 + + def added_source_variables(self) -> List[str]: + """ + Returns + ------- + List of new source variables created by using the subtitles plugin + """ + return ["lang", "subtitles_ext"] + + +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 modify_entry(self, entry: Entry) -> Optional[Entry]: + requested_subtitles = entry.kwargs("requested_subtitles") + languages = sorted(requested_subtitles.keys()) + entry.add_variables( + variables_to_add={ + "subtitles_ext": self.plugin_options.subtitles_type, + "lang": ",".join(languages), + } + ) + + return entry + + def post_process_entry(self, entry: Entry) -> Optional[FileMetadata]: + """ + Creates an entry's NFO file using values defined in the metadata options + + Parameters + ---------- + entry: + Entry to create subtitles for + """ + requested_subtitles = entry.kwargs("requested_subtitles") + file_metadata: Optional[FileMetadata] = None + langs = list(requested_subtitles.keys()) + + if self.plugin_options.embed_subtitles: + file_metadata = FileMetadata(f"Embedded subtitles with lang(s) {', '.join(langs)}") + if self.plugin_options.subtitles_name: + for lang in langs: + subtitle_file_name = f"{entry.uid}.{lang}.{self.plugin_options.subtitles_type}" + output_subtitle_file_name = self.overrides.apply_formatter( + formatter=self.plugin_options.subtitles_name, + entry=entry, + function_overrides={"lang": lang}, + ) + + self.save_file( + file_name=subtitle_file_name, + output_file_name=output_subtitle_file_name, + entry=entry, + ) + + return file_metadata diff --git a/src/ytdl_sub/subscriptions/subscription.py b/src/ytdl_sub/subscriptions/subscription.py index 12aa3b94..d234bfa5 100644 --- a/src/ytdl_sub/subscriptions/subscription.py +++ b/src/ytdl_sub/subscriptions/subscription.py @@ -1,17 +1,16 @@ import contextlib -import copy import os import shutil 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 YTDLOptions @@ -19,7 +18,7 @@ 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 @@ -83,7 +82,7 @@ class Subscription: return self.__preset_options.downloader_options @property - def plugins(self) -> List[Tuple[Type[Plugin], PluginOptions]]: + def plugins(self) -> PresetPlugins: """ Returns ------- @@ -178,6 +177,7 @@ class Subscription: entry=entry, ) + # TODO: see if entry even has a thumbnail if self.output_options.thumbnail_name: output_thumbnail_name = self.overrides.apply_formatter( formatter=self.output_options.thumbnail_name, entry=entry @@ -240,7 +240,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, @@ -262,25 +262,21 @@ 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) - ytdl_options["writethumbnail"] = True - if dry_run: - ytdl_options["skip_download"] = True - ytdl_options["writethumbnail"] = 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 - ) - 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() + with self._subscription_download_context_managers(): downloader = self.downloader_class( download_options=self.downloader_options, enhanced_download_archive=self._enhanced_download_archive, - ytdl_options=ytdl_options, + ytdl_options_builder=ytdl_options_builder, ) for entry in downloader.download(): 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..22b83c02 --- /dev/null +++ b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py @@ -0,0 +1,147 @@ +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 + + @property + def _global_options(self) -> Dict: + """ + Returns + ------- + ytdl-options to apply to every run no matter what + """ + ytdl_options = { + # Download all files in the format of {id}.{ext} + "outtmpl": str(Path(self._working_directory) / "%(id)s.%(ext)s") + } + + if ( + self._downloader.supports_download_archive + and self._preset.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 _dry_run_options(self) -> Dict: + return { + "skip_download": True, + "writethumbnail": False, + # TODO: find a way to not write subtitles; using `simulate: True` breaks tests + } + + @property + def _output_options(self) -> Dict: + ytdl_options = {} + output_options = self._preset.output_options + + if output_options.thumbnail_name: + ytdl_options["writethumbnail"] = True + + return ytdl_options + + @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 {} + + builder = YTDLOptionsBuilder().add({"writesubtitles": True}) + subtitle_options: SubtitleOptions = subtitle_plugin.plugin_options + + if subtitle_options.embed_subtitles: + builder.add( + {"postprocessors": [{"key": "FFmpegEmbedSubtitle", "already_have_subtitle": True}]} + ) + + if subtitle_options.subtitles_name: + builder.add( + { + "postprocessors": [ + { + "key": "FFmpegSubtitlesConvertor", + "format": subtitle_options.subtitles_type, + } + ], + } + ) + + # If neither subtitles_name or embed_subtitles is set, do not set any other flags + if not builder.to_dict(): + return {} + + builder.add( + { + "writeautomaticsub": subtitle_options.allow_auto_generated_subtitles, + "subtitleslangs": subtitle_options.languages, + } + ) + + return builder.to_dict() + + @property + def _user_ytdl_options(self) -> Dict: + return self._preset.ytdl_options.dict + + def builder(self) -> YTDLOptionsBuilder: + """ + Returns + ------- + YTDLOptionsBuilder + Builder with values set based on the subscription + """ + ytdl_options_builder = YTDLOptionsBuilder().add(self._global_options) + if self._dry_run: + ytdl_options_builder.add( + self._subtitle_options, self._user_ytdl_options, self._dry_run_options + ) + else: + ytdl_options_builder.add( + self._output_options, self._subtitle_options, self._user_ytdl_options + ) + + return ytdl_options_builder diff --git a/src/ytdl_sub/validators/strict_dict_validator.py b/src/ytdl_sub/validators/strict_dict_validator.py index f03107c7..e76f4eec 100644 --- a/src/ytdl_sub/validators/strict_dict_validator.py +++ b/src/ytdl_sub/validators/strict_dict_validator.py @@ -27,13 +27,6 @@ class StrictDictValidator(DictValidator): if required_key not in self._dict: raise self._validation_exception(f"missing the required field '{required_key}'") - # Ensure an empty dict was not passed as the value - if not self._dict: - raise self._validation_exception( - f"at least one of the following fields must be defined: " - f"{', '.join(self._optional_keys)}'" - ) - # Ensure all keys are either required or optional keys if no extra field are allowed if not self._allow_extra_keys: for object_key in self._keys: 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/plugins/test_subtitles.py b/tests/e2e/plugins/test_subtitles.py new file mode 100644 index 00000000..bea0cbe5 --- /dev/null +++ b/tests/e2e/plugins/test_subtitles.py @@ -0,0 +1,87 @@ +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_video_subs_embed_preset_dict(output_directory): + return { + "preset": "yt_music_video", + "youtube": {"video_url": "https://www.youtube.com/watch?v=2lAe1cqCOXo"}, + # override the output directory with our fixture-generated dir + "output_options": {"output_directory": output_directory}, + "subtitles": { + "embed_subtitles": True, + "languages": ["en", "de"], + "allow_auto_generated_subtitles": True, + }, + # download the worst format so it is fast + "ytdl_options": { + "format": "worst[ext=mp4]", + }, + "overrides": {"artist": "JMC"}, + } + + +@pytest.fixture +def test_single_video_subs_embed_and_file_preset_dict(single_video_subs_embed_preset_dict): + single_video_subs_embed_preset_dict["subtitles"][ + "subtitles_name" + ] = "{music_video_name}.{lang}.{subtitles_ext}" + return single_video_subs_embed_preset_dict + + +class TestSubtitles: + @pytest.mark.parametrize("dry_run", [True, False]) + def test_subtitles_embedded( + self, + music_video_config, + single_video_subs_embed_preset_dict, + output_directory, + dry_run, + ): + subscription = Subscription.from_dict( + config=music_video_config, + preset_name="subtitles_embedded_test", + preset_dict=single_video_subs_embed_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_subtitles_embedded.txt", + ) + assert_expected_downloads( + output_directory=output_directory, + dry_run=dry_run, + expected_download_summary_file_name="plugins/test_subtitles_embedded.json", + ) + + @pytest.mark.parametrize("dry_run", [True, False]) + def test_subtitles_embedded_and_file( + self, + music_video_config, + test_single_video_subs_embed_and_file_preset_dict, + output_directory, + dry_run, + ): + subscription = Subscription.from_dict( + config=music_video_config, + preset_name="subtitles_embedded_and_file_test", + preset_dict=test_single_video_subs_embed_and_file_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_subtitles_embedded_and_file.txt", + ) + assert_expected_downloads( + output_directory=output_directory, + dry_run=dry_run, + expected_download_summary_file_name="plugins/test_subtitles_embedded_and_file.json", + ) diff --git a/tests/e2e/resources/expected_downloads_summaries/plugins/test_subtitles_embedded.json b/tests/e2e/resources/expected_downloads_summaries/plugins/test_subtitles_embedded.json new file mode 100644 index 00000000..d39fc75f --- /dev/null +++ b/tests/e2e/resources/expected_downloads_summaries/plugins/test_subtitles_embedded.json @@ -0,0 +1,5 @@ +{ + "JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind-thumb.jpg": "704246dd78074e8a0ec001dd8d03fd60", + "JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.mp4": "90ae944811ca3312fcb3175ea32a0aa5", + "JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.nfo": "ffa10f1cbc098ace7b1c7a8fbe3097a8" +} \ No newline at end of file diff --git a/tests/e2e/resources/expected_downloads_summaries/plugins/test_subtitles_embedded_and_file.json b/tests/e2e/resources/expected_downloads_summaries/plugins/test_subtitles_embedded_and_file.json new file mode 100644 index 00000000..1c1c6901 --- /dev/null +++ b/tests/e2e/resources/expected_downloads_summaries/plugins/test_subtitles_embedded_and_file.json @@ -0,0 +1,7 @@ +{ + "JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind-thumb.jpg": "704246dd78074e8a0ec001dd8d03fd60", + "JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.de.srt": "b343c3bb9257b7ee7ba38f570a115b37", + "JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.en.srt": "fe8c6ee92cae6e059fd80fd61691adbe", + "JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.mp4": "90ae944811ca3312fcb3175ea32a0aa5", + "JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.nfo": "ffa10f1cbc098ace7b1c7a8fbe3097a8" +} \ No newline at end of file diff --git a/tests/e2e/resources/transaction_log_summaries/plugins/test_subtitles_embedded.txt b/tests/e2e/resources/transaction_log_summaries/plugins/test_subtitles_embedded.txt new file mode 100644 index 00000000..03daf6ee --- /dev/null +++ b/tests/e2e/resources/transaction_log_summaries/plugins/test_subtitles_embedded.txt @@ -0,0 +1,12 @@ +Files created in '{output_directory}' +---------------------------------------- +JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind-thumb.jpg +JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.mp4 + Embedded subtitles with lang(s) en, de +JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.nfo + NFO tags: + musicvideo: + album: Music Videos + artist: JMC + title: YouTube Rewind 2019: For the Record | #YouTubeRewind + year: 2019 \ No newline at end of file diff --git a/tests/e2e/resources/transaction_log_summaries/plugins/test_subtitles_embedded_and_file.txt b/tests/e2e/resources/transaction_log_summaries/plugins/test_subtitles_embedded_and_file.txt new file mode 100644 index 00000000..d47900c3 --- /dev/null +++ b/tests/e2e/resources/transaction_log_summaries/plugins/test_subtitles_embedded_and_file.txt @@ -0,0 +1,14 @@ +Files created in '{output_directory}' +---------------------------------------- +JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind-thumb.jpg +JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.de.srt +JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.en.srt +JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.mp4 + Embedded subtitles with lang(s) en, de +JMC - YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.nfo + NFO tags: + musicvideo: + album: Music Videos + artist: JMC + title: YouTube Rewind 2019: For the Record | #YouTubeRewind + year: 2019 \ No newline at end of file