diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 2af2b909..082145f1 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -33,6 +33,7 @@ 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.entries.entry_parent import EntryParent +from ytdl_sub.entries.variables.kwargs import COMMENTS from ytdl_sub.entries.variables.kwargs import DOWNLOAD_INDEX from ytdl_sub.entries.variables.kwargs import PLAYLIST_ENTRY from ytdl_sub.entries.variables.kwargs import REQUESTED_SUBTITLES @@ -453,6 +454,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT], ABC): REQUESTED_SUBTITLES: download_entry.kwargs_get(REQUESTED_SUBTITLES), # Same with sponsorblock chapters SPONSORBLOCK_CHAPTERS: download_entry.kwargs_get(SPONSORBLOCK_CHAPTERS), + COMMENTS: download_entry.kwargs_get(COMMENTS), # Tracks number of entries downloaded DOWNLOAD_INDEX: download_idx, # Tracks number of entries with the same upload date to make them unique diff --git a/src/ytdl_sub/entries/variables/kwargs.py b/src/ytdl_sub/entries/variables/kwargs.py index 59e2729e..715ba43a 100644 --- a/src/ytdl_sub/entries/variables/kwargs.py +++ b/src/ytdl_sub/entries/variables/kwargs.py @@ -44,8 +44,10 @@ DOWNLOAD_INDEX = _("download_index", backend=True) UPLOAD_DATE_INDEX = _("upload_date_index", backend=True) REQUESTED_SUBTITLES = _("requested_subtitles", backend=True) CHAPTERS = _("chapters", backend=True) +YTDL_SUB_CUSTOM_CHAPTERS = _("ytdl_sub_custom_chapters", backend=True) SPONSORBLOCK_CHAPTERS = _("sponsorblock_chapters", backend=True) SPLIT_BY_CHAPTERS_PARENT_ENTRY = _("split_by_chapters_parent_entry", backend=True) +COMMENTS = _("comments", backend=True) UID = _("id") EXTRACTOR = _("extractor") EPOCH = _("epoch") diff --git a/src/ytdl_sub/plugins/chapters.py b/src/ytdl_sub/plugins/chapters.py index 1862de82..3c9bc01d 100644 --- a/src/ytdl_sub/plugins/chapters.py +++ b/src/ytdl_sub/plugins/chapters.py @@ -7,6 +7,8 @@ from typing import Set from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder from ytdl_sub.entries.entry import Entry +from ytdl_sub.entries.variables.kwargs import COMMENTS +from ytdl_sub.entries.variables.kwargs import YTDL_SUB_CUSTOM_CHAPTERS from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.plugin import PluginOptions from ytdl_sub.utils.chapters import Chapters @@ -69,7 +71,14 @@ class ChaptersOptions(PluginOptions): presets: my_example_preset: chapters: + # Embedded Chapter Fields embed_chapters: True + allow_chapters_from_comments: False + remove_chapters_regex: + - "Intro" + - "Outro" + + # Sponsorblock Fields sponsorblock_categories: - "outro" - "selfpromo" @@ -79,15 +88,12 @@ class ChaptersOptions(PluginOptions): - "music_offtopic" - "intro" remove_sponsorblock_categories: "all" - remove_chapters_regex: - - "Intro" - - "Outro" force_key_frames: False + """ _optional_keys = { "embed_chapters", - "allow_chapters_from_description", "allow_chapters_from_comments", "embed_chapter_timestamps", "sponsorblock_categories", @@ -116,9 +122,6 @@ class ChaptersOptions(PluginOptions): self._embed_chapter_timestamps = self._validate_key_if_present( "embed_chapter_timestamps", StringValidator ) - self._allow_chapters_from_description = self._validate_key_if_present( - key="allow_chapters_from_description", validator=BoolValidator, default=False - ).value self._allow_chapters_from_comments = self._validate_key_if_present( key="allow_chapters_from_comments", validator=BoolValidator, default=False ).value @@ -128,6 +131,11 @@ class ChaptersOptions(PluginOptions): "Must specify sponsorblock_categories if you are going to remove any of them" ) + if self._remove_sponsorblock_categories and self._allow_chapters_from_comments: + raise self._validation_exception( + "Cannot remove sponsorblock categories and embed chapters from comments" + ) + if self._embed_chapters and self._embed_chapter_timestamps: raise self._validation_exception( "Cannot embed chapters from the source and from a timestamp file" @@ -210,6 +218,14 @@ class ChaptersOptions(PluginOptions): return self._embed_chapter_timestamps.value return None + @property + def allow_chapters_from_comments(self) -> bool: + """ + Optional. If chapters do not exist in the video/description itself, attempt to scrape + comments to find the chapters. Defaults to False. + """ + return self._allow_chapters_from_comments + class ChaptersPlugin(Plugin[ChaptersOptions]): plugin_options_type = ChaptersOptions @@ -251,6 +267,9 @@ class ChaptersPlugin(Plugin[ChaptersOptions]): } ) + if self.plugin_options.allow_chapters_from_comments: + builder.add({"getcomments": True}) + if self._is_removing_chapters: remove_chapters_post_processor = { "key": "ModifyChapters", @@ -312,15 +331,28 @@ class ChaptersPlugin(Plugin[ChaptersOptions]): ------- entry """ + chapters = Chapters.from_empty() + if self.plugin_options.embed_chapter_timestamps and not self.is_dry_run: chapters = Chapters.from_timestamps_file( chapters_file_path=self.plugin_options.embed_chapter_timestamps ) - set_ffmpeg_metadata_chapters( - file_path=entry.get_download_file_path(), - chapters=chapters, - file_duration_sec=entry.kwargs("duration"), - ) + + if not _contains_any_chapters(entry) and self.plugin_options.allow_chapters_from_comments: + for comment in entry.kwargs_get(COMMENTS, []): + chapters = Chapters.from_string(comment.get("text", "")) + if not chapters.is_empty(): + break + + if not chapters.is_empty(): + entry.add_kwargs({YTDL_SUB_CUSTOM_CHAPTERS: chapters.to_file_metadata_dict()}) + + if not self.is_dry_run: + set_ffmpeg_metadata_chapters( + file_path=entry.get_download_file_path(), + chapters=chapters, + file_duration_sec=entry.kwargs("duration"), + ) return entry @@ -335,11 +367,20 @@ class ChaptersPlugin(Plugin[ChaptersOptions]): ------- FileMetadata outlining which chapters/SponsorBlock segments got removed """ - if self.plugin_options.embed_chapter_timestamps: - chapters = Chapters.from_timestamps_file( - chapters_file_path=self.plugin_options.embed_chapter_timestamps + if custom_chapters_metadata := entry.kwargs_get(YTDL_SUB_CUSTOM_CHAPTERS): + title: str = "" + + if self.plugin_options.embed_chapter_timestamps: + title = "Chapters embedded from timestamp file" + elif self.plugin_options.allow_chapters_from_comments: + title = "Chapters from comments" + + assert title, "title should not be empty" + return FileMetadata.from_dict( + value_dict=custom_chapters_metadata, + title=title, + sort_dict=False, # timestamps + titles are already sorted ) - return chapters.to_file_metadata(title="Chapters embedded from timestamp file") if self.plugin_options.embed_chapters: metadata_dict = {} diff --git a/src/ytdl_sub/utils/chapters.py b/src/ytdl_sub/utils/chapters.py index da4f26d9..ace4393d 100644 --- a/src/ytdl_sub/utils/chapters.py +++ b/src/ytdl_sub/utils/chapters.py @@ -2,6 +2,7 @@ import json import os import re import subprocess +from typing import Dict from typing import List from typing import Optional from typing import Tuple @@ -158,6 +159,14 @@ class Chapters: """ return self.timestamps[0].timestamp_sec == 0 + def to_file_metadata_dict(self) -> Dict: + """ + Returns + ------- + Metadata dict + """ + return {ts.readable_str: title for ts, title in zip(self.timestamps, self.titles)} + def to_file_metadata(self, title: Optional[str] = None) -> FileMetadata: """ Parameters @@ -170,7 +179,7 @@ class Chapters: Chapter metadata in the format of { readable_timestamp_str: title } """ return FileMetadata.from_dict( - value_dict={ts.readable_str: title for ts, title in zip(self.timestamps, self.titles)}, + value_dict=self.to_file_metadata_dict(), title=title, sort_dict=False, # timestamps + titles are already sorted ) @@ -319,5 +328,25 @@ class Chapters: return Chapters(timestamps=timestamps, titles=titles) - def __len__(self): + @classmethod + def from_empty(cls) -> "Chapters": + """ + Initialize empty chapters + """ + return Chapters(timestamps=[], titles=[]) + + def __len__(self) -> int: + """ + Returns + ------- + Number of chapters + """ return len(self.timestamps) + + def is_empty(self) -> bool: + """ + Returns + ------- + True if no chapters. False otherwise. + """ + return len(self) == 0 diff --git a/tests/e2e/plugins/test_chapters.py b/tests/e2e/plugins/test_chapters.py index 032f13d5..95149c02 100644 --- a/tests/e2e/plugins/test_chapters.py +++ b/tests/e2e/plugins/test_chapters.py @@ -1,3 +1,5 @@ +from typing import Dict + import pytest from expected_download import assert_expected_downloads from expected_transaction_log import assert_transaction_log_matches @@ -6,7 +8,7 @@ from ytdl_sub.subscriptions.subscription import Subscription @pytest.fixture -def single_video_sponsorblock_and_embedded_subs_preset_dict(output_directory): +def sponsorblock_and_subs_preset_dict(output_directory) -> Dict: return { "preset": "music_video", "download": {"url": "https://www.youtube.com/watch?v=-wJOUAuKZm8"}, @@ -42,19 +44,31 @@ def single_video_sponsorblock_and_embedded_subs_preset_dict(output_directory): } +@pytest.fixture +def chapters_from_comments_preset_dict(sponsorblock_and_subs_preset_dict: Dict) -> Dict: + sponsorblock_and_subs_preset_dict["download"][ + "url" + ] = "https://www.youtube.com/watch?v=MO5AWAqe01Y" + sponsorblock_and_subs_preset_dict["chapters"] = { + "embed_chapters": True, + "allow_chapters_from_comments": True, + } + return sponsorblock_and_subs_preset_dict + + 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, + sponsorblock_and_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, + preset_dict=sponsorblock_and_subs_preset_dict, ) transaction_log = subscription.download(dry_run=dry_run) @@ -76,24 +90,24 @@ class TestChapters: def test_chapters_from_timestamp_file_with_subs( self, music_video_config, - single_video_sponsorblock_and_embedded_subs_preset_dict, + sponsorblock_and_subs_preset_dict, timestamps_file_path, output_directory, dry_run, ): # Test chapters and video tags, throw in a video tag with special chars while we are at it - single_video_sponsorblock_and_embedded_subs_preset_dict["chapters"] = { + sponsorblock_and_subs_preset_dict["chapters"] = { "embed_chapters": False, "embed_chapter_timestamps": timestamps_file_path, } - single_video_sponsorblock_and_embedded_subs_preset_dict["video_tags"] = { + sponsorblock_and_subs_preset_dict["video_tags"] = { "tags": {"description": "🎸 / ' \" \n newline?"} } subscription = Subscription.from_dict( config=music_video_config, preset_name="chapters_from_timestamps_with_subs", - preset_dict=single_video_sponsorblock_and_embedded_subs_preset_dict, + preset_dict=sponsorblock_and_subs_preset_dict, ) transaction_log = subscription.download(dry_run=dry_run) @@ -110,3 +124,30 @@ class TestChapters: "JMC/JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case.mp4" ], ) + + @pytest.mark.parametrize("dry_run", [True, False]) + def test_chapters_from_comments( + self, + music_video_config, + chapters_from_comments_preset_dict, + timestamps_file_path, + output_directory, + dry_run, + ): + subscription = Subscription.from_dict( + config=music_video_config, + preset_name="chapters_from_comments", + preset_dict=chapters_from_comments_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/chapters/test_chapters_from_comments.txt", + ) + assert_expected_downloads( + output_directory=output_directory, + dry_run=dry_run, + expected_download_summary_file_name="plugins/chapters/test_chapters_from_comments.json", + ) diff --git a/tests/resources/expected_downloads_summaries/plugins/chapters/test_chapters_from_comments.json b/tests/resources/expected_downloads_summaries/plugins/chapters/test_chapters_from_comments.json new file mode 100644 index 00000000..4f097093 --- /dev/null +++ b/tests/resources/expected_downloads_summaries/plugins/chapters/test_chapters_from_comments.json @@ -0,0 +1,7 @@ +{ + ".ytdl-sub-chapters_from_comments-download-archive.json": "122723ce8d257eebb05178daa26141f6", + "JMC/JMC - Move 78 - Automated Improvisation [Full Album]-thumb.jpg": "c12e6a6f242680d1096a1a99d74a62c6", + "JMC/JMC - Move 78 - Automated Improvisation [Full Album].info.json": "fc463545b52f7f07265f07df882ca9d7", + "JMC/JMC - Move 78 - Automated Improvisation [Full Album].mp4": "8b8a9a731a19bc37feb091559f97ebbc", + "JMC/JMC - Move 78 - Automated Improvisation [Full Album].nfo": "7a65b184d24c68fc0ec5380432250f5b" +} \ No newline at end of file diff --git a/tests/resources/transaction_log_summaries/plugins/chapters/test_chapters_from_comments.txt b/tests/resources/transaction_log_summaries/plugins/chapters/test_chapters_from_comments.txt new file mode 100644 index 00000000..00325e22 --- /dev/null +++ b/tests/resources/transaction_log_summaries/plugins/chapters/test_chapters_from_comments.txt @@ -0,0 +1,26 @@ +Files created: +---------------------------------------- +{output_directory} + .ytdl-sub-chapters_from_comments-download-archive.json +{output_directory}/JMC + JMC - Move 78 - Automated Improvisation [Full Album]-thumb.jpg + JMC - Move 78 - Automated Improvisation [Full Album].info.json + JMC - Move 78 - Automated Improvisation [Full Album].mp4 + Chapters from comments: + 0:00: 01. The Lonely Tears of Lee Seedol + 4:30: 02. But What If We're Wrong + 9:16: 03. Follow the Earworm Pt.2 + 12:25: 04. Keyword Salad + 16:48: 05. Ultra Natural + 20:47: 06. Flight Instructions + 25:46: 07. Dawn of the Useless Class + 29:58: 08. Schnitzel Whisperer + 32:16: 09. Teilo + Embedded subtitles with lang(s) en, de + JMC - Move 78 - Automated Improvisation [Full Album].nfo + NFO tags: + musicvideo: + album: Music Videos + artist: JMC + title: Move 78 - Automated Improvisation [Full Album] + year: 2022 \ No newline at end of file diff --git a/tests/unit/utils/test_chapters.py b/tests/unit/utils/test_chapters.py index aefab07f..dba87516 100644 --- a/tests/unit/utils/test_chapters.py +++ b/tests/unit/utils/test_chapters.py @@ -71,4 +71,4 @@ class TestChapters: def test_chapters_from_str_2(self, chapter_description_2): chapters = Chapters.from_string(chapter_description_2) assert len(chapters) == 8 - assert chapters.timestamps[1].readable_str == "2:41" \ No newline at end of file + assert chapters.timestamps[1].readable_str == "2:41"