From cb37f5938ce07d1c3802ea699f5475ef17d70b81 Mon Sep 17 00:00:00 2001 From: jbannon Date: Mon, 15 Aug 2022 20:38:54 +0000 Subject: [PATCH] chapters and subs actually working --- src/ytdl_sub/config/preset_class_mappings.py | 2 + src/ytdl_sub/plugins/chapters.py | 147 ++++++++++++++++-- src/ytdl_sub/plugins/split_by_chapters.py | 1 - src/ytdl_sub/plugins/subtitles.py | 37 ++--- .../subscription_ytdl_options.py | 43 ++--- tests/e2e/plugins/test_chapters.py | 71 +++++++++ .../test_chapters_sb_and_embedded_subs.json | 5 + .../test_chapters_sb_and_embedded_subs.txt | 19 +++ 8 files changed, 259 insertions(+), 66 deletions(-) create mode 100644 tests/e2e/plugins/test_chapters.py create mode 100644 tests/e2e/resources/expected_downloads_summaries/plugins/test_chapters_sb_and_embedded_subs.json create mode 100644 tests/e2e/resources/transaction_log_summaries/plugins/test_chapters_sb_and_embedded_subs.txt diff --git a/src/ytdl_sub/config/preset_class_mappings.py b/src/ytdl_sub/config/preset_class_mappings.py index d3bd1386..de94220c 100644 --- a/src/ytdl_sub/config/preset_class_mappings.py +++ b/src/ytdl_sub/config/preset_class_mappings.py @@ -10,6 +10,7 @@ from ytdl_sub.downloaders.youtube.playlist import YoutubePlaylistDownloader from ytdl_sub.downloaders.youtube.split_video import YoutubeSplitVideoDownloader from ytdl_sub.downloaders.youtube.video import YoutubeVideoDownloader from ytdl_sub.plugins.audio_extract import AudioExtractPlugin +from ytdl_sub.plugins.chapters import ChaptersPlugin from ytdl_sub.plugins.music_tags import MusicTagsPlugin from ytdl_sub.plugins.nfo_tags import NfoTagsPlugin from ytdl_sub.plugins.output_directory_nfo_tags import OutputDirectoryNfoTagsPlugin @@ -115,6 +116,7 @@ class PluginMapping: "output_directory_nfo_tags": OutputDirectoryNfoTagsPlugin, "regex": RegexPlugin, "subtitles": SubtitlesPlugin, + "chapters": ChaptersPlugin, } @classmethod diff --git a/src/ytdl_sub/plugins/chapters.py b/src/ytdl_sub/plugins/chapters.py index 7bcb2b0f..0c4e9bc0 100644 --- a/src/ytdl_sub/plugins/chapters.py +++ b/src/ytdl_sub/plugins/chapters.py @@ -1,14 +1,18 @@ +import collections import re 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.regex_validator import RegexListValidator from ytdl_sub.validators.string_select_validator import StringSelectValidator +from ytdl_sub.validators.validators import BoolValidator from ytdl_sub.validators.validators import ListValidator SPONSORBLOCK_HIGHLIGHT_CATEGORIES: Set[str] = {"poi_highlight"} @@ -25,19 +29,22 @@ SPONSORBLOCK_CATEGORIES: Set[str] = SPONSORBLOCK_HIGHLIGHT_CATEGORIES | { class SponsorBlockCategoriesValidator(StringSelectValidator): - _expected_value_type = "sponsorblock category" + _expected_value_type_name = "sponsorblock category" _select_values = {"all"} | SPONSORBLOCK_CATEGORIES class SponsorBlockCategoryListValidator(ListValidator[SponsorBlockCategoriesValidator]): - _expected_value_type = "sponsorblock category" + _expected_value_type_name = "sponsorblock category" _inner_list_type = SponsorBlockCategoriesValidator class ChaptersOptions(PluginOptions): """ - Add chapters to video files if they are present. Options to add SponsorBlock chapters and - remove them or existing chapters. + Add chapters to video files if they are present. Additional options to add SponsorBlock + chapters and remove specific ones. Can also remove chapters using regex patterns. + + Note that at this time, chapter removal with regex will not work with chapters added via + timestamp file. Usage: @@ -56,14 +63,16 @@ class ChaptersOptions(PluginOptions): - "intro" remove_sponsorblock_categories: "all" remove_chapters_regex: - - "sponsor" - + - "Intro" + - "Outro" + force_key_frames: False """ _optional_keys = { "sponsorblock_categories", "remove_sponsorblock_categories", "remove_chapters_regex", + "force_key_frames", } def __init__(self, name, value): @@ -77,17 +86,33 @@ class ChaptersOptions(PluginOptions): self._remove_chapters_regex = self._validate_key_if_present( key="remove_chapters_regex", validator=RegexListValidator ) + self._force_key_frames = self._validate_key_if_present( + key="force_key_frames", validator=BoolValidator, default=False + ).value + + if self._remove_sponsorblock_categories and not self._sponsorblock_categories: + raise self._validation_exception( + "Must specify sponsorblock_categories if you are going to remove any of them" + ) @property def sponsorblock_categories(self) -> Optional[List[str]]: if self._sponsorblock_categories: - return [validator.value for validator in self._sponsorblock_categories.list] + category_list = [validator.value for validator in self._sponsorblock_categories.list] + if "all" in category_list: + return list(SPONSORBLOCK_CATEGORIES) + return category_list return None @property def remove_sponsorblock_categories(self) -> Optional[List[str]]: if self._remove_sponsorblock_categories: - return [validator.value for validator in self._remove_sponsorblock_categories.list] + category_list = [ + validator.value for validator in self._remove_sponsorblock_categories.list + ] + if "all" in category_list: + return list(set(self.sponsorblock_categories) - SPONSORBLOCK_HIGHLIGHT_CATEGORIES) + return category_list return None @property @@ -96,12 +121,108 @@ class ChaptersOptions(PluginOptions): return [validator.compiled_regex for validator in self._remove_chapters_regex.list] return None + @property + def force_key_frames(self) -> bool: + return self._force_key_frames + class ChaptersPlugin(Plugin[ChaptersOptions]): plugin_options_type = ChaptersOptions - def modify_entry(self, entry: Entry) -> Optional[Entry]: - """ - Tags the entry's audio file using values defined in the metadata options - """ - return entry + @property + def _is_removing_chapters(self) -> bool: + return ( + self.plugin_options.remove_chapters_regex is not None + or self.plugin_options.remove_sponsorblock_categories is not None + ) + + def ytdl_options(self) -> Optional[Dict]: + builder = YTDLOptionsBuilder() + if self.plugin_options.sponsorblock_categories: + builder.add( + { + "postprocessors": [ + { + "key": "SponsorBlock", + "when": "pre_process", + "categories": self.plugin_options.sponsorblock_categories, + }, + {"key": "ModifyChapters"}, + ] + } + ) + + # Always add chapters + builder.add({"postprocessors": [{"key": "FFmpegMetadata", "add_chapters": True}]}) + + if self._is_removing_chapters: + remove_chapters_post_processor = { + "key": "ModifyChapters", + "force_keyframes": self.plugin_options.force_key_frames, + } + if self.plugin_options.remove_sponsorblock_categories is not None: + remove_chapters_post_processor[ + "remove_sponsor_segments" + ] = self.plugin_options.remove_sponsorblock_categories + if self.plugin_options.remove_chapters_regex is not None: + remove_chapters_post_processor[ + "remove_chapters_patterns" + ] = self.plugin_options.remove_chapters_regex + + builder.add( + { + "postprocessors": [ + remove_chapters_post_processor, + {"key": "FFmpegMetadata", "add_chapters": True}, # re-add chapters + ] + } + ) + + return builder.to_dict() + + def _chapters(self, entry: Entry) -> List[Dict]: + if entry.kwargs_contains("chapters"): + return entry.kwargs("chapters") + return [] + + def _sponsorblock_chapters(self, entry: Entry) -> List[Dict]: + if entry.kwargs_contains("sponsorblock_chapters"): + return entry.kwargs("sponsorblock_chapters") + return [] + + def _get_removed_chapters(self, entry: Entry) -> List[str]: + removed_chapters: List[str] = [] + for pattern in self.plugin_options.remove_chapters_regex or []: + for chapter in self._chapters(entry): + if pattern.search(chapter["title"]): + removed_chapters.append(chapter["title"]) + return removed_chapters + + def _get_removed_sponsorblock_category_counts(self, entry: Entry) -> Dict: + removed_category_counts = collections.Counter() + for category in self.plugin_options.remove_sponsorblock_categories or []: + for chapter in self._sponsorblock_chapters(entry): + if chapter["category"] == category: + removed_category_counts.update({chapter["title"]: 1}) + + # To make this reproducible, we must sort categories with equal counts by name, + return dict( + sorted( + removed_category_counts.most_common(), + key=lambda name_count: (-name_count[1], name_count[0]), + ) + ) + + def post_process_entry(self, entry: Entry) -> Optional[FileMetadata]: + metadata_dict = {} + removed_chapters = self._get_removed_chapters(entry) + removed_sponsorblock = self._get_removed_sponsorblock_category_counts(entry) + + if removed_chapters: + metadata_dict["Removed Chapter(s)"] = ", ".join(removed_chapters) + if removed_sponsorblock: + metadata_dict["Removed SponsorBlock Category Count(s)"] = removed_sponsorblock + + return FileMetadata.from_dict( + value_dict=metadata_dict, title="Embedded Chapters", sort_dict=False + ) diff --git a/src/ytdl_sub/plugins/split_by_chapters.py b/src/ytdl_sub/plugins/split_by_chapters.py index a608a9db..6fdc8d3c 100644 --- a/src/ytdl_sub/plugins/split_by_chapters.py +++ b/src/ytdl_sub/plugins/split_by_chapters.py @@ -1,4 +1,3 @@ -from typing import Dict from typing import Optional from ytdl_sub.entries.entry import Entry diff --git a/src/ytdl_sub/plugins/subtitles.py b/src/ytdl_sub/plugins/subtitles.py index 1724deb6..ddfc2617 100644 --- a/src/ytdl_sub/plugins/subtitles.py +++ b/src/ytdl_sub/plugins/subtitles.py @@ -135,40 +135,37 @@ 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, - }, - } - ) + builder = YTDLOptionsBuilder().add({"writesubtitles": True}) if self.plugin_options.embed_subtitles: - ytdl_options_builder.add( + builder.add( + {"postprocessors": [{"key": "FFmpegEmbedSubtitle", "already_have_subtitle": True}]} + ) + + if self.plugin_options.subtitles_name: + builder.add( { "postprocessors": [ - # already_have_subtitle=True means keep the subtitle files - {"key": "FFmpegEmbedSubtitle", "already_have_subtitle": write_subtitle_file} - ] + { + "key": "FFmpegSubtitlesConvertor", + "format": self.plugin_options.subtitles_type, + } + ], } ) # If neither subtitles_name or embed_subtitles is set, do not set any other flags - if not ytdl_options_builder.to_dict(): + if not builder.to_dict(): return {} - return ytdl_options_builder.add( + builder.add( { "writeautomaticsub": self.plugin_options.allow_auto_generated_subtitles, "subtitleslangs": self.plugin_options.languages, } - ).to_dict() + ) + + return builder.to_dict() def modify_entry(self, entry: Entry) -> Optional[Entry]: requested_subtitles = entry.kwargs("requested_subtitles") diff --git a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py index e4dc95a3..4d8cca11 100644 --- a/src/ytdl_sub/subscriptions/subscription_ytdl_options.py +++ b/src/ytdl_sub/subscriptions/subscription_ytdl_options.py @@ -9,8 +9,8 @@ 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.audio_extract import AudioExtractPlugin +from ytdl_sub.plugins.chapters import ChaptersPlugin 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 @@ -98,39 +98,14 @@ class SubscriptionYTDLOptions: # TODO: warn here return {} - # TODO: Use subtitle ytdl_options - builder = YTDLOptionsBuilder().add({"writesubtitles": True}) - subtitle_options: SubtitleOptions = subtitle_plugin.plugin_options + return subtitle_plugin.ytdl_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(): + @property + def _chapter_options(self) -> Dict: + if not (chapters_plugin := self._get_plugin(ChaptersPlugin)): return {} - builder.add( - { - "writeautomaticsub": subtitle_options.allow_auto_generated_subtitles, - "subtitleslangs": subtitle_options.languages, - } - ) - - return builder.to_dict() + return chapters_plugin.ytdl_options() @property def _user_ytdl_options(self) -> Dict: @@ -146,12 +121,16 @@ class SubscriptionYTDLOptions: 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 + self._subtitle_options, + self._chapter_options, + self._user_ytdl_options, + self._dry_run_options, ) else: ytdl_options_builder.add( self._output_options, self._subtitle_options, + self._chapter_options, self._audio_extract_options, self._user_ytdl_options, ) diff --git a/tests/e2e/plugins/test_chapters.py b/tests/e2e/plugins/test_chapters.py new file mode 100644 index 00000000..5e122b03 --- /dev/null +++ b/tests/e2e/plugins/test_chapters.py @@ -0,0 +1,71 @@ +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_sponsorblock_and_embedded_subs_preset_dict(output_directory): + return { + "preset": "yt_music_video", + "youtube": {"video_url": "https://www.youtube.com/watch?v=-wJOUAuKZm8"}, + # 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, + }, + "chapters": { + "sponsorblock_categories": [ + "outro", + "selfpromo", + "preview", + "interaction", + "sponsor", + "music_offtopic", + "intro", + ], + "remove_sponsorblock_categories": "all", + "remove_chapters_regex": [ + "Intro", + "Outro", + ], + }, + # download the worst format so it is fast + "ytdl_options": { + "format": "worst[ext=mp4]", + }, + "overrides": {"artist": "JMC"}, + } + + +class TestChapters: + @pytest.mark.parametrize("dry_run", [True, False]) + def test_chapters_sponsorblock_and_removal_with_subs( + self, + music_video_config, + single_video_sponsorblock_and_embedded_subs_preset_dict, + output_directory, + dry_run, + ): + subscription = Subscription.from_dict( + config=music_video_config, + preset_name="sponsorblock_with_embedded_subs_test", + preset_dict=single_video_sponsorblock_and_embedded_subs_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_chapters_sb_and_embedded_subs.txt", + regenerate_transaction_log=True, + ) + assert_expected_downloads( + output_directory=output_directory, + dry_run=dry_run, + expected_download_summary_file_name="plugins/test_chapters_sb_and_embedded_subs.json", + regenerate_expected_download_summary=True, + ) diff --git a/tests/e2e/resources/expected_downloads_summaries/plugins/test_chapters_sb_and_embedded_subs.json b/tests/e2e/resources/expected_downloads_summaries/plugins/test_chapters_sb_and_embedded_subs.json new file mode 100644 index 00000000..264aa7d5 --- /dev/null +++ b/tests/e2e/resources/expected_downloads_summaries/plugins/test_chapters_sb_and_embedded_subs.json @@ -0,0 +1,5 @@ +{ + "JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case-thumb.jpg": "a81457393418b5abed785a82122f3352", + "JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case.mp4": "ff829c47452173eac92e09507feb64a0", + "JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case.nfo": "0c06fe6874588209fccbd9276a446750" +} \ No newline at end of file diff --git a/tests/e2e/resources/transaction_log_summaries/plugins/test_chapters_sb_and_embedded_subs.txt b/tests/e2e/resources/transaction_log_summaries/plugins/test_chapters_sb_and_embedded_subs.txt new file mode 100644 index 00000000..127e6543 --- /dev/null +++ b/tests/e2e/resources/transaction_log_summaries/plugins/test_chapters_sb_and_embedded_subs.txt @@ -0,0 +1,19 @@ +Files created in '{output_directory}' +---------------------------------------- +JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case-thumb.jpg +JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case.mp4 + Embedded Chapters + Removed Chapter(s): Intro, Outro + Removed SponsorBlock Category Count(s): + Sponsor: 2 + Endcards/Credits: 1 + Intermission/Intro Animation: 1 + Unpaid/Self Promotion: 1 + Embedded subtitles with lang(s) en, de +JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case.nfo + NFO tags: + musicvideo: + album: Music Videos + artist: JMC + title: This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case + year: 2021 \ No newline at end of file