From 96c9c4b152e494ade84d4ccdfa9d890447e960ac Mon Sep 17 00:00:00 2001 From: jbannon Date: Mon, 8 Aug 2022 06:10:09 +0000 Subject: [PATCH] Subtitle support WIP --- src/ytdl_sub/config/preset_options.py | 13 ++++++++++ src/ytdl_sub/downloaders/downloader.py | 1 + src/ytdl_sub/entries/entry.py | 12 +++++++++ .../entries/variables/entry_variables.py | 11 ++++++++ src/ytdl_sub/subscriptions/subscription.py | 26 ++++++++++++++++++- tests/e2e/youtube/test_video.py | 2 +- 6 files changed, 63 insertions(+), 2 deletions(-) diff --git a/src/ytdl_sub/config/preset_options.py b/src/ytdl_sub/config/preset_options.py index a7f15ed5..a889022d 100644 --- a/src/ytdl_sub/config/preset_options.py +++ b/src/ytdl_sub/config/preset_options.py @@ -110,6 +110,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", @@ -131,6 +132,9 @@ 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 @@ -174,6 +178,15 @@ 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: """ diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 5f19999a..d9ea96b7 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -319,6 +319,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT] if only_info_json: extract_info_ytdl_options["skip_download"] = True extract_info_ytdl_options["writethumbnail"] = False + extract_info_ytdl_options["writesubtitles"] = False ytdl_options_overrides = dict(ytdl_options_overrides, **extract_info_ytdl_options) diff --git a/src/ytdl_sub/entries/entry.py b/src/ytdl_sub/entries/entry.py index 5970bc26..b440f139 100644 --- a/src/ytdl_sub/entries/entry.py +++ b/src/ytdl_sub/entries/entry.py @@ -58,6 +58,18 @@ class Entry(EntryVariables, BaseEntry): return None + def get_download_subtitles_name(self) -> str: + """ + Returns + ------- + The download subtitle's file name + """ + return f"{self.uid}.{self.subtitles_ext}" + + def get_download_subtitles_path(self) -> str: + """Returns the entry's thumbnail's file path to where it was downloaded""" + return str(Path(self.working_directory()) / self.get_download_subtitles_name()) + @final def is_downloaded(self) -> bool: """ diff --git a/src/ytdl_sub/entries/variables/entry_variables.py b/src/ytdl_sub/entries/variables/entry_variables.py index dfb1b1ca..1c2ff226 100644 --- a/src/ytdl_sub/entries/variables/entry_variables.py +++ b/src/ytdl_sub/entries/variables/entry_variables.py @@ -112,6 +112,17 @@ class EntryVariables(SourceVariables): """ return "jpg" + @property + def subtitles_ext(self: BaseEntry) -> str: + """ + 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. + """ + return "srt" + @property def upload_date(self: BaseEntry) -> str: """ diff --git a/src/ytdl_sub/subscriptions/subscription.py b/src/ytdl_sub/subscriptions/subscription.py index 12aa3b94..acef9971 100644 --- a/src/ytdl_sub/subscriptions/subscription.py +++ b/src/ytdl_sub/subscriptions/subscription.py @@ -178,6 +178,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 @@ -193,6 +194,18 @@ class Subscription: entry=entry, ) + # TODO: see if entry even has subtitles + if self.output_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): """ @@ -266,10 +279,21 @@ class Subscription: # 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 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 diff --git a/tests/e2e/youtube/test_video.py b/tests/e2e/youtube/test_video.py index a569827d..4b65df75 100644 --- a/tests/e2e/youtube/test_video.py +++ b/tests/e2e/youtube/test_video.py @@ -74,7 +74,7 @@ class TestYoutubeVideo: preset_dict=single_video_preset_dict, ) - transaction_log = single_video_subscription.download() + transaction_log = single_video_subscription.download(dry_run=dry_run) assert_transaction_log_matches( output_directory=output_directory, transaction_log=transaction_log,