From 966864c9b4af6d010e8b358990925d761cc79a67 Mon Sep 17 00:00:00 2001 From: jbannon Date: Thu, 23 Jun 2022 06:54:13 +0000 Subject: [PATCH] chapters class, trying to add chapter splitting and general ffmpeg metadata creation --- setup.cfg | 2 +- .../downloaders/youtube/merge_playlist.py | 55 +++++--- .../downloaders/youtube/split_video.py | 66 ++------- src/ytdl_sub/utils/chapters.py | 132 ++++++++++++++++++ src/ytdl_sub/utils/ffmpeg.py | 57 ++++++++ tests/e2e/youtube/test_merge_playlist.py | 6 +- 6 files changed, 241 insertions(+), 77 deletions(-) create mode 100644 src/ytdl_sub/utils/chapters.py diff --git a/setup.cfg b/setup.cfg index 145b8682..51df09d2 100644 --- a/setup.cfg +++ b/setup.cfg @@ -24,13 +24,13 @@ package_dir = packages=find: install_requires = + yt-dlp argparse==1.4.0 dicttoxml==1.7.4 mergedeep==1.3.4 mediafile==0.9.0 Pillow==9.1.0 PyYAML==6.0 - yt-dlp==2022.4.8 [options.packages.find] where=src diff --git a/src/ytdl_sub/downloaders/youtube/merge_playlist.py b/src/ytdl_sub/downloaders/youtube/merge_playlist.py index b61e873e..89ce0b09 100644 --- a/src/ytdl_sub/downloaders/youtube/merge_playlist.py +++ b/src/ytdl_sub/downloaders/youtube/merge_playlist.py @@ -3,18 +3,20 @@ import os.path import re from pathlib import Path from shutil import copyfile +from typing import Any from typing import Dict from typing import List from typing import Tuple -from ytdl_sub.downloaders.youtube_downloader import YoutubeDownloader, \ - YoutubePlaylistDownloaderOptions +from ytdl_sub.downloaders.youtube_downloader import YoutubeDownloader +from ytdl_sub.downloaders.youtube_downloader import YoutubePlaylistDownloaderOptions from ytdl_sub.downloaders.youtube_downloader import YoutubeVideoDownloaderOptions from ytdl_sub.entries.youtube import YoutubePlaylistVideo from ytdl_sub.entries.youtube import YoutubeVideo from ytdl_sub.utils.exceptions import ValidationException from ytdl_sub.utils.ffmpeg import FFMPEG from ytdl_sub.utils.thumbnail import convert_download_thumbnail +from ytdl_sub.validators.validators import BoolValidator from ytdl_sub.validators.validators import StringValidator ############################################################################### @@ -35,7 +37,8 @@ class YoutubeMergePlaylistDownloaderOptions(YoutubePlaylistDownloaderOptions): # required download_strategy: "merge_playlist" playlist_url: "TODO" - chapter_name: "{title}" + # optional + add_chapters: False CLI usage: @@ -62,9 +65,20 @@ class YoutubeMergePlaylistDownloaderOptions(YoutubePlaylistDownloaderOptions): """ _required_keys = {"playlist_url"} + _optional_keys = {"add_chapters"} - # def __init__(self, name, value): - # super().__init__(name, value) + def __init__(self, name, value): + super().__init__(name, value) + self._add_chapters = self._validate_key_if_present( + "add_chapters", validator=BoolValidator, default=False + ).value + + @property + def add_chapters(self) -> bool: + """ + Whether to add chapters using each video's title in the merged playlist. Defaults to false. + """ + return self._add_chapters class YoutubeMergePlaylistDownloader( @@ -85,22 +99,27 @@ class YoutubeMergePlaylistDownloader( """ return dict( super().ytdl_option_defaults(), - **{"postprocessors": [ - { - "key": "FFmpegVideoRemuxer", - "when": "post_process", - "preferedformat": "mkv", - }, - { - "key": "FFmpegConcat", - "when": "playlist", - } - ]}, + **{ + "playlistreverse": True, + "postprocessors": [ + { + "key": "FFmpegVideoRemuxer", + "when": "post_process", + "preferedformat": "mkv", + }, + { + "key": "FFmpegConcat", + "when": "playlist", + }, + ], + }, ) def download(self) -> List[YoutubeVideo]: """Download a single Youtube video, then split it into multiple videos""" - split_videos: List[YoutubePlaylistVideo] = [] entry_dict = self.extract_info(url=self.download_options.playlist_url) - return [] + if self.download_options.add_chapters: + raise NotImplemented("TODO") + + return [YoutubeVideo(entry_dict=entry_dict, working_directory=self.working_directory)] diff --git a/src/ytdl_sub/downloaders/youtube/split_video.py b/src/ytdl_sub/downloaders/youtube/split_video.py index c0be4792..9f34d4ef 100644 --- a/src/ytdl_sub/downloaders/youtube/split_video.py +++ b/src/ytdl_sub/downloaders/youtube/split_video.py @@ -11,6 +11,8 @@ from ytdl_sub.downloaders.youtube_downloader import YoutubeDownloader from ytdl_sub.downloaders.youtube_downloader import YoutubeVideoDownloaderOptions from ytdl_sub.entries.youtube import YoutubePlaylistVideo from ytdl_sub.entries.youtube import YoutubeVideo +from ytdl_sub.utils.chapters import Chapters +from ytdl_sub.utils.chapters import Timestamp from ytdl_sub.utils.exceptions import ValidationException from ytdl_sub.utils.ffmpeg import FFMPEG from ytdl_sub.utils.thumbnail import convert_download_thumbnail @@ -29,56 +31,11 @@ def _split_video_uid(source_uid: str, idx: int) -> str: return f"{source_uid}___{idx}" -def _parse_split_timestamp_file(split_timestamp_path: str) -> Tuple[List[str], List[str]]: - """ - Returns two lists, one containing timestamps in HH:MM:SS format, and the other titles - """ - if not os.path.isfile(split_timestamp_path): - raise ValidationException( - f"split_timestamp file path '{split_timestamp_path}' does not exist." - ) - - with open(split_timestamp_path, "r", encoding="utf-8") as file: - lines = file.readlines() - - timestamps: List[str] = [] - titles: List[str] = [] - idx = 0 - for idx, line in enumerate(lines): - match = _SPLIT_TIMESTAMP_REGEX.match(line) - if not match: - break - - timestamp = match.group(1) - title = match.group(2) - match len(timestamp): - case 4: # 0:00 - timestamp = f"00:0{timestamp}" - case 5: # 00:00 - timestamp = f"00:{timestamp}" - case 7: # 0:00:00 - timestamp = f"0{timestamp}" - case _: - pass - - assert len(timestamp) == 8 - timestamps.append(timestamp) - titles.append(title) - - if idx not in (len(lines) - 1, len(lines) - 2): - raise ValidationException( - f"split_timestamp file '{split_timestamp_path} is not formatted correctly. " - f"Each line must be formatted as '0:00 title' - a timestamp, space, then title." - ) - - return timestamps, titles - - def _split_video_ffmpeg_cmd( - input_file: str, output_file: str, timestamps: List[str], idx: int + input_file: str, output_file: str, timestamps: List[Timestamp], idx: int ) -> List[str]: - timestamp_begin = timestamps[idx] - timestamp_end = timestamps[idx + 1] if idx + 1 < len(timestamps) else "" + timestamp_begin = timestamps[idx].timestamp_str + timestamp_end = timestamps[idx + 1].timestamp_str if idx + 1 < len(timestamps) else "" cmd = ["-i", input_file, "-ss", timestamp_begin] if timestamp_end: @@ -192,9 +149,7 @@ class YoutubeSplitVideoDownloader( """Download a single Youtube video, then split it into multiple videos""" split_videos: List[YoutubePlaylistVideo] = [] - timestamps, titles = _parse_split_timestamp_file( - split_timestamp_path=self.download_options.split_timestamps - ) + chapters = Chapters.from_file(chapters_file_path=self.download_options.split_timestamps) entry_dict = self.extract_info(url=self.download_options.video_url) entry = YoutubeVideo(entry_dict=entry_dict, working_directory=self.working_directory) @@ -202,7 +157,7 @@ class YoutubeSplitVideoDownloader( # when copying it convert_download_thumbnail(entry=entry) - for idx, title in enumerate(titles): + for idx, title in enumerate(chapters.titles): new_uid = _split_video_uid(source_uid=entry.uid, idx=idx) # Get the input/output file paths @@ -215,7 +170,10 @@ class YoutubeSplitVideoDownloader( # Run ffmpeg to create the split the video FFMPEG.run( _split_video_ffmpeg_cmd( - input_file=input_file, output_file=output_file, timestamps=timestamps, idx=idx + input_file=input_file, + output_file=output_file, + timestamps=chapters.timestamps, + idx=idx, ) ) # Copy the thumbnail @@ -227,7 +185,7 @@ class YoutubeSplitVideoDownloader( source_entry_dict=entry_dict, title=title, idx=idx, - split_video_count=len(timestamps), + split_video_count=len(chapters.timestamps), ) ) diff --git a/src/ytdl_sub/utils/chapters.py b/src/ytdl_sub/utils/chapters.py new file mode 100644 index 00000000..8a4b7b35 --- /dev/null +++ b/src/ytdl_sub/utils/chapters.py @@ -0,0 +1,132 @@ +import os +import re +from typing import List + +from ytdl_sub.utils.exceptions import ValidationException + + +class Timestamp: + + # Captures the following formats: + # 0:00 title + # 00:00 title + # 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)$") + + @classmethod + def _normalize_timestamp_str(cls, timestamp_str: str) -> str: + match = cls._SPLIT_TIMESTAMP_REGEX.match(timestamp_str) + if not match: + raise ValueError(f"Cannot parse youtube timestamp '{timestamp_str}'") + + timestamp = match.group(1) + match len(timestamp): + case 4: # 0:00 + timestamp = f"00:0{timestamp}" + case 5: # 00:00 + timestamp = f"00:{timestamp}" + case 7: # 0:00:00 + timestamp = f"0{timestamp}" + case _: + pass + + assert len(timestamp) == 8 + return timestamp + + def __init__(self, timestamp_sec: int): + self._timestamp_sec = timestamp_sec + + @property + def timestamp_sec(self) -> int: + return self._timestamp_sec + + @property + def timestamp_str(self) -> str: + seconds = self.timestamp_sec + + hours = int(seconds / 3600) + seconds -= hours * 3600 + + minutes = int(seconds / 60) + seconds -= minutes * 60 + + return f"{str(hours).zfill(2)}:{str(minutes).zfill(2)}:{str(seconds).zfill(2)}" + + @classmethod + def from_seconds(cls, timestamp_sec: int) -> "Timestamp": + """ + Parameters + ---------- + timestamp_sec + Timestamp in number of seconds + """ + return cls(timestamp_sec=timestamp_sec) + + @classmethod + def from_str(cls, timestamp_str: str) -> "Timestamp": + """ + Parameters + ---------- + timestamp_str + Timestamp in the form of "HH:MM:SS" + """ + hour_minute_second = cls._normalize_timestamp_str(timestamp_str).split(":") + if len(hour_minute_second) != 3: + raise ValueError("Youtube timestamp must be in the form of 'HH:MM:SS'") + + hour, minute, second = tuple(x for x in hour_minute_second) + try: + return cls(timestamp_sec=(int(hour) * 3600) + (int(minute) * 60) + int(second)) + except ValueError as cast_exception: + raise ValueError( + "Youtube timestamp must be in the form of 'HH:MM:SS'" + ) from cast_exception + + +class Chapters: + """ + Represents a list of (timestamps, titles) + """ + + def __init__(self, timestamps: List[Timestamp], titles: List[str], duration: Timestamp): + self.timestamps = timestamps + self.titles = titles + self.duration = duration + + def contains_zero_timestamp(self) -> bool: + return self.timestamps[0].timestamp_sec == 0 + + @classmethod + def from_file(cls, chapters_file_path: str, duration: Timestamp) -> "Chapters": + 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.lstrip().split(maxsplit=1) + + # Allow the last line to be blank + if idx == len(lines) - 1 and not line.strip(): + break + + if len(line_split) != 2: + raise ValueError( + f"Chapter/Timestamp file line could not be parsed: '{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, duration=duration) diff --git a/src/ytdl_sub/utils/ffmpeg.py b/src/ytdl_sub/utils/ffmpeg.py index 0539d2c9..842b0234 100644 --- a/src/ytdl_sub/utils/ffmpeg.py +++ b/src/ytdl_sub/utils/ffmpeg.py @@ -1,11 +1,25 @@ import subprocess +import tempfile +from typing import Dict from typing import List +from typing import Optional +from ytdl_sub.utils.chapters import Chapters from ytdl_sub.utils.exceptions import ValidationException from ytdl_sub.utils.logger import Logger logger = Logger.get(name="ffmpeg") +_FFMPEG_METADATA_SPECIAL_CHARS = ["=", ";", "#", "\n", "\\"] + + +def _ffmpeg_metadata_escape(str_to_escape: str) -> str: + # backslash at the end of the list is intentional + for special_char in _FFMPEG_METADATA_SPECIAL_CHARS: + str_to_escape.replace(special_char, f"\\{special_char}") + + return str_to_escape + class FFMPEG: @classmethod @@ -33,3 +47,46 @@ class FFMPEG: cmd.extend(ffmpeg_args) logger.debug("Running %s", " ".join(cmd)) subprocess.run(cmd, check=True) + + +def add_metadata( + file_path: str, + metadata: Optional[Dict[str, str]], + chapters: Optional[Chapters], +) -> None: + + if not metadata and not chapters: + return + + lines = [";FFMETADATA1"] + + if metadata: + for key, value in metadata.items(): + lines.append(f"{_ffmpeg_metadata_escape(key)}={_ffmpeg_metadata_escape(value)}") + + if chapters: + if not chapters.contains_zero_timestamp(): + raise ValueError("Chapters must contain a zero timestamp") + + for idx in range(len(chapters.timestamps)): + start = chapters.timestamps[idx].timestamp_sec + end = ( + chapters.timestamps[idx + 1].timestamp_sec + if idx < len(chapters.timestamps) - 1 + else chapters.duration + ) + + lines.append("") + lines.append("[CHAPTER]") + lines.append("TIMEBASE=1") + lines.append(f"START={start}") + lines.append(f"END={end}") + lines.append(f"title={_ffmpeg_metadata_escape(chapters.titles[idx])}") + + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", encoding="utf-8") as tmp_file: + tmp_file.writelines(lines) + tmp_file.flush() + + yield tmp_file.name + + return None diff --git a/tests/e2e/youtube/test_merge_playlist.py b/tests/e2e/youtube/test_merge_playlist.py index d9e7a950..bee2e842 100644 --- a/tests/e2e/youtube/test_merge_playlist.py +++ b/tests/e2e/youtube/test_merge_playlist.py @@ -32,13 +32,13 @@ def subscription_dict(output_directory, subscription_name): "preset": "yt_music_video_playlist", "youtube": { "download_strategy": "merge_playlist", - "playlist_url": "https://youtube.com/playlist?list=PL5BC0FC26BECA5A35" + "playlist_url": "https://youtube.com/playlist?list=PL5BC0FC26BECA5A35", }, # override the output directory with our fixture-generated dir "output_options": {"output_directory": output_directory}, # download the worst format so it is fast "ytdl_options": { - "format": "worst", + "format": "best", }, "overrides": {"artist": "JMC"}, } @@ -88,7 +88,6 @@ def expected_playlist_download(): # fmt: on - class TestYoutubeMergePlaylist: """ Downloads my old minecraft youtube channel, pretends they are music videos. Ensure the above @@ -100,4 +99,3 @@ class TestYoutubeMergePlaylist: ): playlist_subscription.download() expected_playlist_download.assert_files_exist(relative_directory=output_directory) -