chapters and subs actually working

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

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
+    Add chapters to video files if they are present. Additional options to add SponsorBlock
-    remove them or existing chapters.
+    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]:
+    @property
-        """
+    def _is_removing_chapters(self) -> bool:
-        Tags the entry's audio file using values defined in the metadata options
+        return (
-        """
+            self.plugin_options.remove_chapters_regex is not None
-        return entry
+            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
+        )

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

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()
+        builder = YTDLOptionsBuilder().add({"writesubtitles": True})
-
-        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(
+            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")

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
+        return subtitle_plugin.ytdl_options()
-        builder = YTDLOptionsBuilder().add({"writesubtitles": True})
-        subtitle_options: SubtitleOptions = subtitle_plugin.plugin_options
 
-        if subtitle_options.embed_subtitles:
+    @property
-            builder.add(
+    def _chapter_options(self) -> Dict:
-                {"postprocessors": [{"key": "FFmpegEmbedSubtitle", "already_have_subtitle": True}]}
+        if not (chapters_plugin := self._get_plugin(ChaptersPlugin)):
-            )
-
-        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(
+        return chapters_plugin.ytdl_options()
-            {
-                "writeautomaticsub": subtitle_options.allow_auto_generated_subtitles,
-                "subtitleslangs": subtitle_options.languages,
-            }
-        )
-
-        return builder.to_dict()
 
     @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,
             )

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,
+        )

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"
+}

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