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

[FEATURE] Chapters plugin with SponsorBlock support (#178)

Browse source
This commit is contained in:
Jesse Bannon 2022-08-15 17:01:03 -07:00 committed by GitHub
parent 871c051aa8
commit 2a666ef70b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 440 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
docs/config.rst
View file

@ -177,6 +177,14 @@ audio_extract
-------------------------------------------------------------------------------
chapters
''''''''
.. autoclass:: ytdl_sub.plugins.chapters.ChaptersOptions()
:members:
:member-order: bysource
-------------------------------------------------------------------------------
music_tags
''''''''''
.. autoclass:: ytdl_sub.plugins.music_tags.MusicTagsOptions()

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

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

1
src/ytdl_sub/downloaders/downloader.py
View file

@ -57,6 +57,7 @@ class Downloader(DownloadArchiver, Generic[DownloaderOptionsT, DownloaderEntryT]
supports_download_archive: bool = True
supports_subtitles: bool = True
supports_chapters: bool = True
_extract_entry_num_retries: int = 5
_extract_entry_retry_wait_sec: int = 3

1
src/ytdl_sub/downloaders/soundcloud/albums_and_singles.py
View file

@ -52,6 +52,7 @@ class SoundcloudAlbumsAndSinglesDownloader(
):
downloader_options_type = SoundcloudAlbumsAndSinglesDownloadOptions
supports_subtitles = False
supports_chapters = False
@classmethod
def ytdl_option_defaults(cls) -> Dict:

1
src/ytdl_sub/downloaders/youtube/merge_playlist.py
View file

@ -65,6 +65,7 @@ class YoutubeMergePlaylistDownloader(
downloader_entry_type = YoutubeVideo
supports_download_archive = False
supports_subtitles = False
supports_chapters = False
@classmethod
def ytdl_option_defaults(cls) -> Dict:

284
src/ytdl_sub/plugins/chapters.py Normal file
View file

@ -0,0 +1,284 @@
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"}
SPONSORBLOCK_CATEGORIES: Set[str] = SPONSORBLOCK_HIGHLIGHT_CATEGORIES | {
"sponsor",
"intro",
"outro",
"selfpromo",
"preview",
"filler",
"interaction",
"music_offtopic",
}
def _chapters(entry: Entry) -> List[Dict]:
if entry.kwargs_contains("chapters"):
return entry.kwargs("chapters")
return []
def _sponsorblock_chapters(entry: Entry) -> List[Dict]:
if entry.kwargs_contains("sponsorblock_chapters"):
return entry.kwargs("sponsorblock_chapters")
return []
class SponsorBlockCategoriesValidator(StringSelectValidator):
_expected_value_type_name = "sponsorblock category"
_select_values = {"all"} | SPONSORBLOCK_CATEGORIES
class SponsorBlockCategoryListValidator(ListValidator[SponsorBlockCategoriesValidator]):
_expected_value_type_name = "sponsorblock category"
_inner_list_type = SponsorBlockCategoriesValidator
class ChaptersOptions(PluginOptions):
"""
Embeds chapters to video files if they are present. Additional options to add SponsorBlock
chapters and remove specific ones. Can also remove chapters using regex.
Note that at this time, chapter removal with regex will not work with chapters added via
timestamp file.
Usage:
.. code-block:: yaml
presets:
my_example_preset:
chapters:
embed_chapters: True
sponsorblock_categories:
- "outro"
- "selfpromo"
- "preview"
- "interaction"
- "sponsor"
- "music_offtopic"
- "intro"
remove_sponsorblock_categories: "all"
remove_chapters_regex:
- "Intro"
- "Outro"
force_key_frames: False
"""
_optional_keys = {
"embed_chapters",
"sponsorblock_categories",
"remove_sponsorblock_categories",
"remove_chapters_regex",
"force_key_frames",
}
def __init__(self, name, value):
super().__init__(name, value)
self._embed_chapters = self._validate_key_if_present(
key="embed_chapters", validator=BoolValidator, default=True
).value
self._sponsorblock_categories = self._validate_key_if_present(
key="sponsorblock_categories", validator=SponsorBlockCategoryListValidator
)
self._remove_sponsorblock_categories = self._validate_key_if_present(
key="remove_sponsorblock_categories", validator=SponsorBlockCategoryListValidator
)
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 embed_chapters(self) -> Optional[bool]:
"""
Optional. Embed chapters into the file. Defaults to True.
"""
return self._embed_chapters
@property
def sponsorblock_categories(self) -> Optional[List[str]]:
"""
Optional. List of SponsorBlock categories to embed as chapters. Supports "sponsor",
"intro", "outro", "selfpromo", "preview", "filler", "interaction", "music_offtopic",
"poi_highlight", or "all" to include all categories.
"""
if self._sponsorblock_categories:
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]]:
"""
Optional. List of SponsorBlock categories to remove from the output file. Can only remove
categories that are specified in ``sponsorblock_categories`` or "all", which removes
everything specified in ``sponsorblock_categories``.
"""
if self._remove_sponsorblock_categories:
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
def remove_chapters_regex(self) -> Optional[List[re.Pattern]]:
"""
Optional. List of regex patterns to match chapter titles against and remove them from the
entry.
"""
if self._remove_chapters_regex:
return [validator.compiled_regex for validator in self._remove_chapters_regex.list]
return None
@property
def force_key_frames(self) -> bool:
"""
Optional. Force keyframes at cuts when removing sections. This is slow due to needing a
re-encode, but the resulting video may have fewer artifacts around the cuts. Defaults to
False.
"""
return self._force_key_frames
class ChaptersPlugin(Plugin[ChaptersOptions]):
plugin_options_type = ChaptersOptions
@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]:
"""
Returns
-------
YTDL options to embed chapters, add/remove SponsorBlock segments, remove chapters via regex
"""
builder = YTDLOptionsBuilder()
if self.plugin_options.sponsorblock_categories:
builder.add(
{
"postprocessors": [
{
"key": "SponsorBlock",
"when": "pre_process",
"categories": self.plugin_options.sponsorblock_categories,
},
{"key": "ModifyChapters"},
]
}
)
if self.plugin_options.embed_chapters:
builder.add(
{
"postprocessors": [
{"key": "FFmpegMetadata", "add_chapters": True, "add_metadata": False}
]
}
)
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
if self.plugin_options.embed_chapters:
builder.add(
{
# re-add chapters
"postprocessors": [
remove_chapters_post_processor,
{"key": "FFmpegMetadata", "add_chapters": True, "add_metadata": False},
]
}
)
return builder.to_dict()
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 _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 _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]:
"""
Parameters
----------
entry:
Entry with possibly removed chapters
Returns
-------
FileMetadata outlining which chapters/SponsorBlock segments got removed
"""
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
)

37
src/ytdl_sub/plugins/subtitles.py
View file

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

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

@ -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,18 @@ 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,
}
)
if not self._downloader.supports_chapters:
# TODO: warn here
return {}
return builder.to_dict()
return chapters_plugin.ytdl_options()
@property
def _user_ytdl_options(self) -> Dict:
@ -146,12 +125,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,
)

9
src/ytdl_sub/validators/regex_validator.py
View file

@ -28,6 +28,15 @@ class RegexValidator(StringValidator):
"""
return self._compiled_regex.groups
@property
def compiled_regex(self) -> re.Pattern:
"""
Returns
-------
The regex compiled
"""
return self._compiled_regex
def match(self, input_str: str) -> Optional[List[str]]:
"""
Parameters

6
tests/e2e/plugins/test_audio_extract.py
View file

@ -15,6 +15,7 @@ def single_song_preset_dict(output_directory):
# download the worst format so it is fast
"ytdl_options": {
"format": "worst[ext=mp4]",
"postprocessor_args": {"ffmpeg": ["-bitexact"]}, # Must add this for reproducibility
},
}
@ -29,6 +30,7 @@ def multiple_songs_preset_dict(output_directory):
# download the worst format so it is fast
"ytdl_options": {
"format": "worst[ext=mp4]",
"postprocessor_args": {"ffmpeg": ["-bitexact"]}, # Must add this for reproducibility
},
}
@ -53,13 +55,11 @@ class TestAudioExtract:
output_directory=output_directory,
transaction_log=transaction_log,
transaction_log_summary_file_name="plugins/test_audio_extract_single.txt",
regenerate_transaction_log=True,
)
assert_expected_downloads(
output_directory=output_directory,
dry_run=dry_run,
expected_download_summary_file_name="plugins/test_audio_extract_single.json",
regenerate_expected_download_summary=True,
)
@pytest.mark.parametrize("dry_run", [True, False])
@ -81,11 +81,9 @@ class TestAudioExtract:
output_directory=output_directory,
transaction_log=transaction_log,
transaction_log_summary_file_name="plugins/test_audio_extract_playlist.txt",
regenerate_transaction_log=True,
)
assert_expected_downloads(
output_directory=output_directory,
dry_run=dry_run,
expected_download_summary_file_name="plugins/test_audio_extract_playlist.json",
regenerate_expected_download_summary=True,
)

73
tests/e2e/plugins/test_chapters.py Normal file
View file

@ -0,0 +1,73 @@
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]",
"postprocessor_args": {"ffmpeg": ["-bitexact"]}, # Must add this for reproducibility
},
"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",
)
assert_expected_downloads(
output_directory=output_directory,
dry_run=dry_run,
expected_download_summary_file_name="plugins/test_chapters_sb_and_embedded_subs.json",
ignore_md5_hashes_for=[
"JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case.mp4"
],
)

6
tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_playlist.json
View file

@ -1,5 +1,5 @@
{
"Jesse's Minecraft Server [Trailer - Feb.1].ogg": "67b19b495756fb6dc332f49bfa2957d3",
"Jesse's Minecraft Server [Trailer - Feb.27].ogg": "273c96652e198171ee60345051addc7a",
"Jesse's Minecraft Server [Trailer - Mar.21].ogg": "e91db58c153f0724aaed5ce53de2a736"
"Jesse's Minecraft Server [Trailer - Feb.1].ogg": "c37c289bc9b7c79464aa8bcc6df423e3",
"Jesse's Minecraft Server [Trailer - Feb.27].ogg": "37def0736bb5c0a7fba0c1685e90bf3c",
"Jesse's Minecraft Server [Trailer - Mar.21].ogg": "1fcc317ee5ce24675f6e4e6c2eae40d3"
}

2
tests/e2e/resources/expected_downloads_summaries/plugins/test_audio_extract_single.json
View file

@ -1,3 +1,3 @@
{
"YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.mp3": "c4ce74e2ddc4e6a0ab823b5b622e2275"
"YouTube Rewind 2019： For the Record ｜ #YouTubeRewind.mp3": "3a156b122bd79c956cce5079d3530cc3"
}

5
tests/e2e/resources/expected_downloads_summaries/plugins/test_chapters_sb_and_embedded_subs.json Normal file
View file

@ -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": "2c417b31c9f5eb8cbecf0bf1fc1f9e53",
"JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case.nfo": "0c06fe6874588209fccbd9276a446750"
}

19
tests/e2e/resources/transaction_log_summaries/plugins/test_chapters_sb_and_embedded_subs.txt Normal file
View file

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