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

[FEATURE] Add support to scrape chapters from comments (#347)

Browse source 
* [FEATURE] Add support to scrape chapters from comments

* tested for comments

* remove timestamps file

* change order of docstrings
This commit is contained in:
Jesse Bannon 2022-11-21 11:19:17 -08:00 committed by GitHub
parent 70c78dc109
commit ffaf07d628
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 233 additions and 115 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/ytdl_sub/downloaders/downloader.py
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

2
src/ytdl_sub/entries/variables/kwargs.py
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")

117
src/ytdl_sub/plugins/chapters.py
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
@ -16,7 +18,6 @@ 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 | {
@ -62,9 +63,6 @@ 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
@ -72,7 +70,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"
@ -82,15 +87,13 @@ class ChaptersOptions(PluginOptions):
- "music_offtopic"
- "intro"
remove_sponsorblock_categories: "all"
remove_chapters_regex:
- "Intro"
- "Outro"
force_key_frames: False
"""
_optional_keys = {
"embed_chapters",
"embed_chapter_timestamps",
"allow_chapters_from_comments",
"sponsorblock_categories",
"remove_sponsorblock_categories",
"remove_chapters_regex",
@ -114,18 +117,18 @@ class ChaptersOptions(PluginOptions):
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
)
self._allow_chapters_from_comments = self._validate_key_if_present(
key="allow_chapters_from_comments", 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"
)
if self._embed_chapters and self._embed_chapter_timestamps:
if self._remove_sponsorblock_categories and self._allow_chapters_from_comments:
raise self._validation_exception(
"Cannot embed chapters from the source and from a timestamp file"
"Cannot remove sponsorblock categories and embed chapters from comments"
)
@property
@ -135,6 +138,24 @@ class ChaptersOptions(PluginOptions):
"""
return self._embed_chapters
@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
@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 sponsorblock_categories(self) -> Optional[List[str]]:
"""
@ -165,16 +186,6 @@ class ChaptersOptions(PluginOptions):
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:
"""
@ -184,27 +195,6 @@ class ChaptersOptions(PluginOptions):
"""
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
@ -246,6 +236,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",
@ -307,15 +300,23 @@ class ChaptersPlugin(Plugin[ChaptersOptions]):
-------
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"),
)
chapters = Chapters.from_empty()
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
@ -330,11 +331,13 @@ 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 = "Chapters from comments"
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 = {}

98
src/ytdl_sub/utils/chapters.py
View file

@ -1,13 +1,12 @@
import json
import os
import re
import subprocess
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from ytdl_sub.entries.entry import Entry
from ytdl_sub.utils.exceptions import ValidationException
from ytdl_sub.utils.file_handler import FileMetadata
@ -19,11 +18,11 @@ class Timestamp:
# 1:00:00 title
# 01:00:00 title
# where capture group 1 and 2 are the timestamp and title, respectively
_SPLIT_TIMESTAMP_REGEX = re.compile(r"^((?:\d\d:)?(?:\d:)?(?:\d)?\d:\d\d)$")
TIMESTAMP_REGEX = re.compile(r"((?:\d\d:)?(?:\d:)?(?:\d)?\d:\d\d)")
@classmethod
def _normalize_timestamp_str(cls, timestamp_str: str) -> str:
match = cls._SPLIT_TIMESTAMP_REGEX.match(timestamp_str)
match = cls.TIMESTAMP_REGEX.match(timestamp_str)
if not match:
raise ValueError(f"Cannot parse youtube timestamp '{timestamp_str}'")
@ -158,6 +157,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,54 +177,48 @@ 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
)
@classmethod
def from_timestamps_file(cls, chapters_file_path: str) -> "Chapters":
def from_string(cls, input_str: str) -> "Chapters":
"""
From a string (description or comment), try to extract Chapters.
The scraping logic is simple, if three or more successive lines have timestamps, grab
as many in succession as possible. Remove the timestamp portion to get the chapter title.
Parameters
----------
chapters_file_path
Path to file containing chapters
input_str
String to scrape
Raises
------
ValidationException
File path does not exist or contains invalid formatting
Returns
-------
Chapters
Could be empty
"""
if not os.path.isfile(chapters_file_path):
raise ValidationException(
f"chapter/timestamp file path '{chapters_file_path}' does not exist."
)
with open(chapters_file_path, "r", encoding="utf-8") as file:
lines = file.readlines()
timestamps: List[Timestamp] = []
titles: List[str] = []
for idx, line in enumerate(lines):
line_split = line.strip().split(maxsplit=1)
for line in input_str.split("\n"):
# Timestamp captured, store it
if match := Timestamp.TIMESTAMP_REGEX.search(line):
timestamp_str = match.group(1)
timestamps.append(Timestamp.from_str(timestamp_str))
# Allow the last line to be blank
if idx == len(lines) - 1 and not line.strip():
break
# Remove timestamp and surrounding whitespace from it
title_str = re.sub(f"\\s*{re.escape(timestamp_str)}\\s*", " ", line).strip()
titles.append(title_str)
elif len(timestamps) >= 3:
return Chapters(timestamps=timestamps, titles=titles)
# Timestamp was not stored, if only contained 1, reset
else:
timestamps = []
titles = []
if len(line_split) != 2:
raise ValidationException(
f"Chapter/Timestamp file '{chapters_file_path}' could not parse '{line}': "
f"must be in the format of 'HH:MM:SS title"
)
timestamp_str, title = tuple(x for x in line_split)
timestamps.append(Timestamp.from_str(timestamp_str))
titles.append(title)
return cls(timestamps=timestamps, titles=titles)
return Chapters(timestamps=timestamps, titles=titles)
@classmethod
def from_embedded_chapters(cls, file_path: str) -> "Chapters":
@ -280,3 +281,26 @@ class Chapters:
titles.append(chapter["title"])
return Chapters(timestamps=timestamps, titles=titles)
@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

44
tests/e2e/plugins/test_chapters.py
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)
@ -73,40 +87,28 @@ class TestChapters:
)
@pytest.mark.parametrize("dry_run", [True, False])
def test_chapters_from_timestamp_file_with_subs(
def test_chapters_from_comments(
self,
music_video_config,
single_video_sponsorblock_and_embedded_subs_preset_dict,
chapters_from_comments_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"] = {
"embed_chapters": False,
"embed_chapter_timestamps": timestamps_file_path,
}
single_video_sponsorblock_and_embedded_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_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/test_chapters_from_ts_with_subs.txt",
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/test_chapters_from_ts_with_subs.json",
ignore_md5_hashes_for=[
"JMC/JMC - This GPU SLIDES into this Case! - Silverstone SUGO 16 ITX Case.mp4"
],
expected_download_summary_file_name="plugins/chapters/test_chapters_from_comments.json",
)

7
tests/resources/expected_downloads_summaries/plugins/chapters/test_chapters_from_comments.json Normal file
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"
}

26
tests/resources/transaction_log_summaries/plugins/chapters/test_chapters_from_comments.txt Normal file
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

52
tests/unit/utils/test_chapters.py
View file

@ -1,5 +1,6 @@
import pytest
from ytdl_sub.utils.chapters import Chapters
from ytdl_sub.utils.chapters import Timestamp
@ -20,3 +21,54 @@ class TestTimestamp:
def test_timestamp_from_str(self, timestamp_str, timestamp_int):
ts = Timestamp.from_str(timestamp_str=timestamp_str)
assert ts.timestamp_sec == timestamp_int
@pytest.fixture
def chapter_description_1() -> str:
return """Support the artist by purchasing the record here:
https://levanika.bandcamp.com/album/p...
Album Art by: po
Tracklist:
00:00 intro
00:58 audioclip p94568
01:18 saskipao
03:32 fifqeby
05:40 dream
07:22 rulji
08:54 modymody
11:23 .ishos
16:09 pos
18:47 lamazybalaxy
Denivarlevy Socials:
https://open.spotify.com/artist/2XdIT...
https://youtube.com/channel/UCgI_vAC3...
"""
@pytest.fixture
def chapter_description_2() -> str:
return """01. 00:00 Ocean
02. 02:41 Dreams
03. 05:16 Future Tales
04. 08:50 Mind Travelling
05. 11:05 Love Supreme
06. 14:17 Reflections
07. 16:32 Moonlight Fading
08. 19:40 Between Two Worlds
"""
class TestChapters:
def test_chapters_from_str_1(self, chapter_description_1):
chapters = Chapters.from_string(chapter_description_1)
assert len(chapters) == 10
assert chapters.timestamps[-1].readable_str == "18:47"
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"