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

chapters and subs actually working

Browse source
This commit is contained in:
jbannon 2022-08-15 20:38:54 +00:00
parent 653681040d
commit cb37f5938c
8 changed files with 259 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

147
src/ytdl_sub/plugins/chapters.py
View file

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

1
src/ytdl_sub/plugins/split_by_chapters.py
View file

@ -1,4 +1,3 @@
from typing import Dict
from typing import Optional
from ytdl_sub.entries.entry import Entry

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

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

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

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

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": "ff829c47452173eac92e09507feb64a0",
"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