tested for comments

This commit is contained in:
Jesse Bannon 2022-11-21 10:36:42 -08:00
parent 69ab4d3e46
commit 61287af12d
8 changed files with 174 additions and 26 deletions

View file

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

View file

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

View file

@ -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 = {}

View file

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

View file

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

View file

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

View file

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