358 lines
13 KiB
Python
358 lines
13 KiB
Python
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.chapters import Chapters
|
|
from ytdl_sub.utils.ffmpeg import set_ffmpeg_metadata_chapters
|
|
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
|
|
from ytdl_sub.validators.validators import StringValidator
|
|
|
|
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") or []
|
|
return []
|
|
|
|
|
|
def _sponsorblock_chapters(entry: Entry) -> List[Dict]:
|
|
if entry.kwargs_contains("sponsorblock_chapters"):
|
|
return entry.kwargs("sponsorblock_chapters") or []
|
|
return []
|
|
|
|
|
|
def _contains_any_chapters(entry: Entry) -> bool:
|
|
return len(_chapters(entry)) > 0 or len(_sponsorblock_chapters(entry)) > 0
|
|
|
|
|
|
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",
|
|
"embed_chapter_timestamps",
|
|
"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
|
|
self._embed_chapter_timestamps = self._validate_key_if_present(
|
|
"embed_chapter_timestamps", StringValidator
|
|
)
|
|
|
|
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"
|
|
)
|
|
|
|
if self._embed_chapters and self._embed_chapter_timestamps:
|
|
raise self._validation_exception(
|
|
"Cannot embed chapters from the source and from a timestamp file"
|
|
)
|
|
|
|
@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
|
|
|
|
@property
|
|
def embed_chapter_timestamps(self) -> Optional[str]:
|
|
"""
|
|
Optional. The path to the file containing the timestamps to embed into the file as
|
|
chapters. Should be formatted as:
|
|
|
|
.. code-block:: markdown
|
|
|
|
0:00 Intro
|
|
0:24 Blackwater Park
|
|
10:23 Bleak
|
|
16:39 Jokes
|
|
1:02:23 Ending
|
|
|
|
This should only be used with single entity download strategies. Otherwise, an entire
|
|
playlist or channel would all the same embedded chapters.
|
|
"""
|
|
if self._embed_chapter_timestamps:
|
|
return self._embed_chapter_timestamps.value
|
|
return None
|
|
|
|
|
|
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 modify_entry(self, entry: Entry) -> Entry:
|
|
"""
|
|
Parameters
|
|
----------
|
|
entry
|
|
Entry to add custom chapters using timestamps if present
|
|
|
|
Returns
|
|
-------
|
|
entry
|
|
"""
|
|
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"),
|
|
)
|
|
|
|
return entry
|
|
|
|
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
|
|
"""
|
|
if self.plugin_options.embed_chapter_timestamps:
|
|
chapters = Chapters.from_timestamps_file(
|
|
chapters_file_path=self.plugin_options.embed_chapter_timestamps
|
|
)
|
|
return chapters.to_file_metadata(title="Chapters embedded from timestamp file")
|
|
|
|
if self.plugin_options.embed_chapters:
|
|
metadata_dict = {}
|
|
removed_chapters = self._get_removed_chapters(entry)
|
|
removed_sponsorblock = self._get_removed_sponsorblock_category_counts(entry)
|
|
|
|
# If no chapters are on the entry, do not report any embedded chapters
|
|
if not _contains_any_chapters(entry):
|
|
return None
|
|
|
|
if removed_chapters:
|
|
metadata_dict["Removed Chapter(s)"] = ", ".join(removed_chapters)
|
|
if removed_sponsorblock:
|
|
metadata_dict["Removed SponsorBlock Category Count(s)"] = removed_sponsorblock
|
|
|
|
# TODO: check if file actually has embedded chapters
|
|
return FileMetadata.from_dict(
|
|
value_dict=metadata_dict, title="Embedded Chapters", sort_dict=False
|
|
)
|
|
|
|
return None
|