From fd88d066f39ecab3704386c97cd68e21839c118d Mon Sep 17 00:00:00 2001 From: jbannon Date: Fri, 24 Jun 2022 06:02:46 +0000 Subject: [PATCH] so very close --- src/ytdl_sub/downloaders/downloader.py | 2 + .../downloaders/youtube/merge_playlist.py | 44 +++++++- src/ytdl_sub/subscriptions/subscription.py | 22 +++- src/ytdl_sub/utils/ffmpeg.py | 106 ++++++++++++------ src/ytdl_sub/utils/thumbnail.py | 4 +- tests/e2e/youtube/test_merge_playlist.py | 25 +---- 6 files changed, 139 insertions(+), 64 deletions(-) diff --git a/src/ytdl_sub/downloaders/downloader.py b/src/ytdl_sub/downloaders/downloader.py index 5bacf31a..6e9f3e66 100644 --- a/src/ytdl_sub/downloaders/downloader.py +++ b/src/ytdl_sub/downloaders/downloader.py @@ -44,6 +44,8 @@ class Downloader(Generic[DownloaderOptionsT, DownloaderEntryT], ABC): downloader_options_type: Type[DownloaderValidator] = DownloaderValidator downloader_entry_type: Type[Entry] = Entry + supports_download_archive: bool = True + @classmethod def ytdl_option_overrides(cls) -> Dict: """Global overrides that even overwrite user input""" diff --git a/src/ytdl_sub/downloaders/youtube/merge_playlist.py b/src/ytdl_sub/downloaders/youtube/merge_playlist.py index ce9c1a0c..3bb82829 100644 --- a/src/ytdl_sub/downloaders/youtube/merge_playlist.py +++ b/src/ytdl_sub/downloaders/youtube/merge_playlist.py @@ -4,6 +4,9 @@ from typing import List from ytdl_sub.downloaders.youtube_downloader import YoutubeDownloader from ytdl_sub.downloaders.youtube_downloader import YoutubePlaylistDownloaderOptions 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.ffmpeg import add_ffmpeg_metadata from ytdl_sub.validators.validators import BoolValidator ############################################################################### @@ -73,6 +76,7 @@ class YoutubeMergePlaylistDownloader( ): downloader_options_type = YoutubeMergePlaylistDownloaderOptions downloader_entry_type = YoutubeVideo + supports_download_archive = False @classmethod def ytdl_option_defaults(cls) -> Dict: @@ -102,11 +106,45 @@ class YoutubeMergePlaylistDownloader( }, ) + def _add_chapters(self, merged_video: YoutubeVideo) -> None: + titles: List[str] = [] + timestamps: List[Timestamp] = [] + + current_timestamp_sec = 0 + for video_entry in merged_video.kwargs("entries"): + timestamps.append(Timestamp(current_timestamp_sec)) + titles.append(video_entry["title"]) + + current_timestamp_sec += video_entry["duration"] + + add_ffmpeg_metadata( + file_path=merged_video.get_download_file_path(), + chapters=Chapters(timestamps=timestamps, titles=titles), + file_duration_sec=merged_video.kwargs("duration"), + ) + + def _to_merged_video(self, entry_dict: Dict) -> YoutubeVideo: + """ + Adds a few entries not included in a playlist entry to make it look like a merged video + entry_dict + """ + # Set the upload date to be the latest playlist video date + entry_dict["upload_date"] = max( + playlist_entry["upload_date"] for playlist_entry in entry_dict["entries"] + ) + entry_dict["duration"] = sum( + playlist_entry["duration"] for playlist_entry in entry_dict["entries"] + ) + entry_dict["ext"] = entry_dict["requested_downloads"][0]["ext"] + return YoutubeVideo(entry_dict=entry_dict, working_directory=self.working_directory) + def download(self) -> List[YoutubeVideo]: """Download a single Youtube video, then split it into multiple videos""" - entry_dict = self.extract_info(url=self.download_options.playlist_url) + merged_video = self._to_merged_video( + entry_dict=self.extract_info(url=self.download_options.playlist_url) + ) if self.download_options.add_chapters: - raise NotImplemented("TODO") + self._add_chapters(merged_video=merged_video) - return [YoutubeVideo(entry_dict=entry_dict, working_directory=self.working_directory)] + return [merged_video] diff --git a/src/ytdl_sub/subscriptions/subscription.py b/src/ytdl_sub/subscriptions/subscription.py index f7a0157e..9d07ffa3 100644 --- a/src/ytdl_sub/subscriptions/subscription.py +++ b/src/ytdl_sub/subscriptions/subscription.py @@ -132,6 +132,18 @@ class Subscription: """ return self.overrides.apply_formatter(formatter=self.output_options.output_directory) + @property + def maintain_download_archive(self) -> bool: + """ + Returns + ------- + Whether to maintain a download archive + """ + return ( + self.output_options.maintain_download_archive + and self.downloader_class.supports_download_archive + ) + def _copy_file_to_output_directory( self, entry: Entry, source_file_path: str, output_file_name: str ): @@ -152,7 +164,7 @@ class Subscription: os.makedirs(os.path.dirname(destination_file_path), exist_ok=True) copyfile(source_file_path, destination_file_path) - if self.output_options.maintain_download_archive: + if self.maintain_download_archive: self._enhanced_download_archive.mapping.add_entry(entry, output_file_name) def _copy_entry_files_to_output_directory(self, entry: Entry): @@ -207,14 +219,14 @@ class Subscription: """ Context manager to initialize the enhanced download archive """ - if self.output_options.maintain_download_archive: + if self.maintain_download_archive: self._enhanced_download_archive.prepare_download_archive() yield # If output options maintains stale file deletion, perform the delete here prior to saving # the download archive - if self.output_options.maintain_download_archive: + if self.maintain_download_archive: date_range_to_keep = self.output_options.get_upload_date_range_to_keep() if date_range_to_keep: self._enhanced_download_archive.remove_stale_files(date_range=date_range_to_keep) @@ -234,7 +246,7 @@ class Subscription: output_directory=self.output_directory, overrides=self.overrides, enhanced_download_archive=self._enhanced_download_archive - if self.output_options.maintain_download_archive + if self.maintain_download_archive else None, ) @@ -253,7 +265,7 @@ class Subscription: download_options=self.downloader_options, ytdl_options=self.ytdl_options.dict, download_archive_file_name=self._enhanced_download_archive.archive_file_name - if self.output_options.maintain_download_archive + if self.maintain_download_archive else None, ) diff --git a/src/ytdl_sub/utils/ffmpeg.py b/src/ytdl_sub/utils/ffmpeg.py index 842b0234..3ed4ea82 100644 --- a/src/ytdl_sub/utils/ffmpeg.py +++ b/src/ytdl_sub/utils/ffmpeg.py @@ -1,5 +1,7 @@ +import shutil import subprocess import tempfile +import termios from typing import Dict from typing import List from typing import Optional @@ -49,44 +51,82 @@ class FFMPEG: subprocess.run(cmd, check=True) -def add_metadata( - file_path: str, - metadata: Optional[Dict[str, str]], - chapters: Optional[Chapters], +def _create_metadata_chapter_entry(start_sec: int, end_sec: int, title: str) -> List[str]: + return [ + "", + "[CHAPTER]", + "TIMEBASE=1/1000", + f"START={start_sec * 1000}", + f"END={end_sec * 1000}", + f"title={_ffmpeg_metadata_escape(title)}", + ] + + +def _create_metadata_chapters(chapters: Chapters, file_duration_sec: int) -> List[str]: + lines: List[str] = [] + + if not chapters.contains_zero_timestamp(): + lines += _create_metadata_chapter_entry( + start_sec=0, + end_sec=chapters.timestamps[0].timestamp_sec, + title="Intro", # TODO: make this configurable + ) + + for idx in range(len(chapters.timestamps) - 1): + lines += _create_metadata_chapter_entry( + start_sec=chapters.timestamps[idx].timestamp_sec, + end_sec=chapters.timestamps[idx + 1].timestamp_sec, + title=chapters.titles[idx], + ) + + # Add the last chapter using the file duration + lines += _create_metadata_chapter_entry( + start_sec=chapters.timestamps[-1].timestamp_sec, + end_sec=file_duration_sec, + title=chapters.titles[-1], + ) + + return lines + + +def add_ffmpeg_metadata( + file_path: str, chapters: Optional[Chapters], file_duration_sec: int ) -> None: + """ + Adds ffmetadata to a file. TODO: support more than just chapters - if not metadata and not chapters: - return - + Parameters + ---------- + file_path + Full path to the file to add metadata to + chapters + Chapters to embed in the file. If a chapter for 0:00 does not exist, one is created + file_duration_sec + Length of the file in seconds + """ 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") + lines += _create_metadata_chapters(chapters=chapters, file_duration_sec=file_duration_sec) - 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 - ) + file_path_ext = file_path.split(".")[-1] + output_file_path = f"{file_path}.out.{file_path_ext}" + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", encoding="utf-8") as metadata_file: + metadata_file.write("\n".join(lines)) + metadata_file.flush() - 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])}") + FFMPEG.run( + [ + "-i", + file_path, + "-i", + metadata_file.name, + "-map_metadata", + "1", + "-codec", + "copy", + output_file_path, + ] + ) - 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 + shutil.move(src=output_file_path, dst=file_path) diff --git a/src/ytdl_sub/utils/thumbnail.py b/src/ytdl_sub/utils/thumbnail.py index 7ac7cb14..c137a1f1 100644 --- a/src/ytdl_sub/utils/thumbnail.py +++ b/src/ytdl_sub/utils/thumbnail.py @@ -10,8 +10,8 @@ from ytdl_sub.entries.entry import Entry def _get_downloaded_thumbnail_path(entry: Entry) -> Optional[str]: - thumbnails = entry.kwargs("thumbnails") - possible_thumbnail_exts = set() + thumbnails = entry.kwargs("thumbnails") or [] + possible_thumbnail_exts = {"jpg", "webp"} # Always check for jpg and webp thumbs # The source `thumbnail` value and the actual downloaded thumbnail extension sometimes do # not match. Find all possible extensions by checking all available thumbnails. diff --git a/tests/e2e/youtube/test_merge_playlist.py b/tests/e2e/youtube/test_merge_playlist.py index bee2e842..bde43d6b 100644 --- a/tests/e2e/youtube/test_merge_playlist.py +++ b/tests/e2e/youtube/test_merge_playlist.py @@ -1,11 +1,8 @@ from pathlib import Path -import mergedeep import pytest -from conftest import assert_debug_log from e2e.expected_download import ExpectedDownload -import ytdl_sub.downloaders.downloader from ytdl_sub.config.config_file import ConfigFile from ytdl_sub.config.preset import Preset from ytdl_sub.subscriptions.subscription import Subscription @@ -33,6 +30,7 @@ def subscription_dict(output_directory, subscription_name): "youtube": { "download_strategy": "merge_playlist", "playlist_url": "https://youtube.com/playlist?list=PL5BC0FC26BECA5A35", + "add_chapters": True, }, # override the output directory with our fixture-generated dir "output_options": {"output_directory": output_directory}, @@ -64,28 +62,13 @@ def playlist_subscription(config, subscription_name, subscription_dict): @pytest.fixture def expected_playlist_download(): - # turn off black formatter here for readability - # fmt: off return ExpectedDownload( expected_md5_file_hashes={ - # Download mapping - Path(".ytdl-sub-jmc-download-archive.json"): "7541aa75606b86bff5ff276895520cf0", - - # Entry files - Path("JMC - Jesse's Minecraft Server [Trailer - Feb.1].jpg"): "048a19cf0f674437351872c3f312ebf1", - Path("JMC - Jesse's Minecraft Server [Trailer - Feb.1].mp4"): "e66287b9832277b6a4d1554e29d9fdcc", - Path("JMC - Jesse's Minecraft Server [Trailer - Feb.1].nfo"): "3d272fe58487b6011ad049b6000b046f", - - Path("JMC - Given to Fly.jpg"): "2e58e4d5f06ce5d1c3336fa493470135", - Path("JMC - Given to Fly.mp4"): "04ab5cb3cc12325d0c96a7cd04a8b91d", - Path("JMC - Given to Fly.nfo"): "0dc578bf5f1ceb6e069a57d329894f35", - - Path("JMC - Indifference (Remastered).jpg"): "9baaddc6b62f5b9ae3781eb4eef0e3b3", - Path("JMC - Indifference (Remastered).mp4"): "025de6099a5c98e6397153c7a62d517d", - Path("JMC - Indifference (Remastered).nfo"): "061b86d9dc8fb39d39feab3292dafeb0", + Path("JMC - Jesse's Minecraft Server.jpg"): "348e3007fc590d0b1e2f6682501b0b5f", + Path("JMC - Jesse's Minecraft Server.mkv"): "a568758afb79f42c7de0af027a4bf62d", + Path("JMC - Jesse's Minecraft Server.nfo"): "10df5dcdb65ab18ecf21b3503c77e48b", } ) - # fmt: on class TestYoutubeMergePlaylist: