danvics/ytdl-sub
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

subtitles now a plugin

Browse source
This commit is contained in:
jbannon 2022-08-11 06:36:10 +00:00
parent baae5ff72d
commit ad17bb6ddb
8 changed files with 248 additions and 127 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
src/ytdl_sub/config/preset.py
View file

@ -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

2
src/ytdl_sub/config/preset_class_mappings.py
View file

@ -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

93
src/ytdl_sub/config/preset_options.py
View file

@ -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

3
src/ytdl_sub/entries/variables/entry_variables.py
View file

@ -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"

8
src/ytdl_sub/plugins/plugin.py
View file

@ -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]:
"""

170
src/ytdl_sub/plugins/subtitles.py Normal file
View file

@ -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

36
src/ytdl_sub/subscriptions/subscription.py
View file

@ -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,

22
src/ytdl_sub/subscriptions/subscription_ytdl_options.py
View file

@ -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: