From 83951aeea05f18bcbed95ec36d7dd169dcfdbb86 Mon Sep 17 00:00:00 2001 From: jbannon Date: Thu, 18 Aug 2022 07:29:05 +0000 Subject: [PATCH] lots of work ahead --- examples/youtube_extract_and_tag_audio.yaml | 24 ++++-- src/ytdl_sub/config/preset_class_mappings.py | 2 + .../downloaders/youtube/split_video.py | 4 +- src/ytdl_sub/downloaders/youtube/video.py | 4 +- src/ytdl_sub/plugins/plugin.py | 20 +++++ src/ytdl_sub/plugins/split_by_chapters.py | 83 ++++++++++++------- src/ytdl_sub/subscriptions/subscription.py | 74 ++++++++++++----- src/ytdl_sub/utils/chapters.py | 29 ++++--- tests/e2e/plugins/test_split_by_chapters.py | 48 +++++++++++ 9 files changed, 219 insertions(+), 69 deletions(-) create mode 100644 tests/e2e/plugins/test_split_by_chapters.py diff --git a/examples/youtube_extract_and_tag_audio.yaml b/examples/youtube_extract_and_tag_audio.yaml index e738ba59..ab81c4f0 100644 --- a/examples/youtube_extract_and_tag_audio.yaml +++ b/examples/youtube_extract_and_tag_audio.yaml @@ -28,11 +28,25 @@ presets: overrides: music_directory: "/path/to/music" - yt_song_playlist: - preset: yt_song - youtube: - download_strategy: "playlist" +# TODO: make a playlist of individual songs into an album. Need playlist_title +# yt_album_as_playlist: +# preset: yt_song +# youtube: +# download_strategy: "playlist" +# +# music_tags: +# tags: +# album: "{playlist_title}" +# track: "{playlist_index}" + + yt_album_as_chapters: + preset: "yt_song" + + split_by_chapters: + when_no_chapters: "pass" music_tags: tags: - track: "{playlist_index}" \ No newline at end of file + title: "{chapter_title}" + album: "{title}" + track: "{chapter_index}" \ No newline at end of file diff --git a/src/ytdl_sub/config/preset_class_mappings.py b/src/ytdl_sub/config/preset_class_mappings.py index de94220c..9554340f 100644 --- a/src/ytdl_sub/config/preset_class_mappings.py +++ b/src/ytdl_sub/config/preset_class_mappings.py @@ -16,6 +16,7 @@ from ytdl_sub.plugins.nfo_tags import NfoTagsPlugin from ytdl_sub.plugins.output_directory_nfo_tags import OutputDirectoryNfoTagsPlugin from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.plugins.regex import RegexPlugin +from ytdl_sub.plugins.split_by_chapters import SplitByChaptersPlugin from ytdl_sub.plugins.subtitles import SubtitlesPlugin from ytdl_sub.plugins.video_tags import VideoTagsPlugin @@ -117,6 +118,7 @@ class PluginMapping: "regex": RegexPlugin, "subtitles": SubtitlesPlugin, "chapters": ChaptersPlugin, + "split_by_chapters": SplitByChaptersPlugin, } @classmethod diff --git a/src/ytdl_sub/downloaders/youtube/split_video.py b/src/ytdl_sub/downloaders/youtube/split_video.py index 7a0f6512..1fe5e284 100644 --- a/src/ytdl_sub/downloaders/youtube/split_video.py +++ b/src/ytdl_sub/downloaders/youtube/split_video.py @@ -145,7 +145,9 @@ class YoutubeSplitVideoDownloader( """Download a single Youtube video, then split it into multiple videos""" split_videos_and_metadata: List[Tuple[YoutubePlaylistVideo, FileMetadata]] = [] - chapters = Chapters.from_timestamps_file(chapters_file_path=self.download_options.split_timestamps) + chapters = Chapters.from_timestamps_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) diff --git a/src/ytdl_sub/downloaders/youtube/video.py b/src/ytdl_sub/downloaders/youtube/video.py index 27b91061..e9a00826 100644 --- a/src/ytdl_sub/downloaders/youtube/video.py +++ b/src/ytdl_sub/downloaders/youtube/video.py @@ -103,7 +103,9 @@ class YoutubeVideoDownloader(YoutubeDownloader[YoutubeVideoDownloaderOptions, Yo return [video] # Otherwise, add the chapters and return the video + chapter metadata - chapters = Chapters.from_timestamps_file(chapters_file_path=self.download_options.chapter_timestamps) + chapters = Chapters.from_timestamps_file( + chapters_file_path=self.download_options.chapter_timestamps + ) if not self.is_dry_run: set_ffmpeg_metadata_chapters( file_path=video.get_download_file_path(), diff --git a/src/ytdl_sub/plugins/plugin.py b/src/ytdl_sub/plugins/plugin.py index 654957e0..640b5226 100644 --- a/src/ytdl_sub/plugins/plugin.py +++ b/src/ytdl_sub/plugins/plugin.py @@ -3,6 +3,7 @@ from typing import Dict from typing import Generic from typing import List from typing import Optional +from typing import Tuple from typing import Type from typing import TypeVar from typing import final @@ -61,6 +62,9 @@ class Plugin(DownloadArchiver, Generic[PluginOptionsT], ABC): plugin_options_type: Type[PluginOptionsT] = NotImplemented + # If the plugin creates multile entries from a single entry + is_split_plugin: bool = False + @final def __init__( self, @@ -81,6 +85,22 @@ class Plugin(DownloadArchiver, Generic[PluginOptionsT], ABC): ytdl options to enable/disable when downloading entries for this specific plugin """ + def split(self, entry: Entry) -> List[Tuple[Entry, FileMetadata]]: + """ + Very specialized function that takes an entry and creates multiple entries from it. + Should mark ``is_split_plugin`` on the plugin class. + + Parameters + ---------- + entry + Entry to create multiple entries from + + Returns + ------- + List of entries and metadata created from the source entry + """ + raise NotImplemented() + # pylint: disable=no-self-use def modify_entry(self, entry: Entry) -> Optional[Entry]: """ diff --git a/src/ytdl_sub/plugins/split_by_chapters.py b/src/ytdl_sub/plugins/split_by_chapters.py index ee913833..1b82f018 100644 --- a/src/ytdl_sub/plugins/split_by_chapters.py +++ b/src/ytdl_sub/plugins/split_by_chapters.py @@ -1,16 +1,31 @@ import copy from pathlib import Path -from typing import Optional, List +from typing import List +from typing import Optional +from typing import Tuple 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, Timestamp +from ytdl_sub.utils.chapters import Chapters +from ytdl_sub.utils.chapters import Timestamp from ytdl_sub.utils.ffmpeg import FFMPEG from ytdl_sub.utils.file_handler import FileHandler +from ytdl_sub.utils.file_handler import FileMetadata from ytdl_sub.utils.thumbnail import convert_download_thumbnail from ytdl_sub.validators.string_select_validator import StringSelectValidator +# +# modify_entry BEFORE SPLIT +# - audio_extract +# - subtitles? +# +# TODO: make regex's modify_entry into a new function +# and call modify_entry before split +# +# maybe modify_downloaded_entry ?? + + def _split_video_ffmpeg_cmd( input_file: str, output_file: str, timestamps: List[Timestamp], idx: int ) -> List[str]: @@ -29,14 +44,14 @@ def _split_video_uid(source_uid: str, idx: int) -> str: class WhenNoChaptersValidator(StringSelectValidator): - _expected_value_type = "when no chapters option" + _expected_value_type_name = "when no chapters option" _select_values = {"pass", "drop", "error"} class SplitByChaptersOptions(PluginOptions): """ - Splits a file by chapters into multiple files. Each file becomes its own entry with ``title`` - set to its chapter name, and is processed separately by other plugins. + Splits a file by chapters into multiple files. Each file becomes its own entry with the + new source variables ``chapter_title``, ``chapter_index``, ``chapter_count``. Usage: @@ -56,6 +71,9 @@ class SplitByChaptersOptions(PluginOptions): key="when_no_chapters", validator=WhenNoChaptersValidator ).value + def added_source_variables(self) -> List[str]: + return ["chapter_title", "chapter_index", "chapter_count"] + @property def when_no_chapters(self) -> str: """ @@ -67,46 +85,53 @@ class SplitByChaptersOptions(PluginOptions): class SplitByChaptersPlugin(Plugin[SplitByChaptersOptions]): plugin_options_type = SplitByChaptersOptions + is_split_plugin = True def _create_split_entry( self, source_entry: Entry, title: str, idx: int, chapters: Chapters - ) -> Entry: + ) -> Tuple[Entry, FileMetadata]: """ Runs ffmpeg to create the split video """ - entry_dict = copy.deepcopy(source_entry) - entry_dict["title"] = title - entry_dict["playlist_index"] = idx + 1 - entry_dict["playlist_count"] = len(chapters.timestamps) - entry_dict["id"] = _split_video_uid(source_uid=entry_dict["id"], idx=idx) + entry = copy.deepcopy(source_entry) - # Remove track and artist since its now split - if "track" in entry_dict: - del entry_dict["track"] - if "artist" in entry_dict: - del entry_dict["artist"] + entry.add_variables( + { + "chapter_title": title, + "chapter_index": idx + 1, + "chapter_count": len(chapters.timestamps), + } + ) + entry._kwargs["id"] = _split_video_uid(source_uid=entry.uid, idx=idx) timestamp_begin = chapters.timestamps[idx].readable_str - timestamp_end = Timestamp(source_entry_dict["duration"]).readable_str + timestamp_end = Timestamp(entry.kwargs("duration")).readable_str if idx + 1 < len(chapters.timestamps): timestamp_end = chapters.timestamps[idx + 1].readable_str - metadata = FileMetadata(metadata=f"{timestamp_begin} - {timestamp_end}") - return ( - YoutubePlaylistVideo(entry_dict=entry_dict, working_directory=self.working_directory), - metadata, + metadata = FileMetadata.from_dict( + value_dict={ + "Split Chapter Source": entry.title, + "Segment": f"{timestamp_begin} - {timestamp_end}", + }, + sort_dict=False, ) - def modify_entry(self, entry: Entry) -> Optional[Entry]: + return entry, metadata + + def split(self, entry: Entry) -> Optional[List[Tuple[Entry, FileMetadata]]]: """ Tags the entry's audio file using values defined in the metadata options """ - split_entries: List[Entry] = [] - chapters = Chapters.from_embedded_chapters(file_path=entry.get_download_file_path()) + split_videos_and_metadata: List[Tuple[Entry, FileMetadata]] = [] - # convert the entry thumbnail early so we do not have to guess the thumbnail extension - # when copying it - if not self.is_dry_run: + if self.is_dry_run: + chapters = None + else: + chapters = Chapters.from_embedded_chapters(file_path=entry.get_download_file_path()) + + # convert the entry thumbnail early so we do not have to guess the thumbnail extension + # when copying it convert_download_thumbnail(entry=entry) for idx, title in enumerate(chapters.titles): @@ -136,8 +161,8 @@ class SplitByChaptersPlugin(Plugin[SplitByChaptersOptions]): # Format the split video as a YoutubePlaylistVideo split_videos_and_metadata.append( - self._create_split_video_entry( - source_entry_dict=entry_dict, title=title, idx=idx, chapters=chapters + self._create_split_entry( + source_entry=entry, title=title, idx=idx, chapters=chapters ) ) diff --git a/src/ytdl_sub/subscriptions/subscription.py b/src/ytdl_sub/subscriptions/subscription.py index d234bfa5..ccd0f969 100644 --- a/src/ytdl_sub/subscriptions/subscription.py +++ b/src/ytdl_sub/subscriptions/subscription.py @@ -19,12 +19,23 @@ from ytdl_sub.downloaders.downloader import DownloaderValidator from ytdl_sub.entries.entry import Entry from ytdl_sub.plugins.plugin import Plugin from ytdl_sub.subscriptions.subscription_ytdl_options import SubscriptionYTDLOptions +from ytdl_sub.utils.exceptions import ValidationException from ytdl_sub.utils.file_handler import FileHandlerTransactionLog from ytdl_sub.utils.file_handler import FileMetadata from ytdl_sub.utils.thumbnail import convert_download_thumbnail from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive +def _get_split_plugin(plugins: List[Plugin]) -> Optional[Plugin]: + split_plugins = [plugin for plugin in plugins if plugin.is_split_plugin] + + if len(split_plugins) == 1: + return split_plugins[0] + if len(split_plugins) > 1: + raise ValidationException("Can not use more than one split plugins at a time") + return None + + class Subscription: """ Subscription classes are the 'controllers' that perform... @@ -251,6 +262,38 @@ class Subscription: return plugins + def _process_entry( + self, plugins: List[Plugin], dry_run: bool, entry: Entry, entry_metadata: FileMetadata + ) -> None: + # First, modify the entry with all plugins + for plugin in plugins: + # Return if it is None, it is indicated to not process any further + if (entry := plugin.modify_entry(entry)) is None: + return + + # Post-process the entry with all plugins + for plugin in plugins: + optional_plugin_entry_metadata = plugin.post_process_entry(entry) + if optional_plugin_entry_metadata: + entry_metadata.extend(optional_plugin_entry_metadata) + + # Then, move it to the output directory + self._move_entry_files_to_output_directory( + dry_run=dry_run, entry=entry, entry_metadata=entry_metadata + ) + + # Re-save the download archive after each entry is moved to the output directory + if self.maintain_download_archive: + self._enhanced_download_archive.save_download_mappings() + + def _process_split_entry( + self, split_plugin: Plugin, plugins: List[Plugin], dry_run: bool, entry: Entry + ) -> None: + for entry, entry_metadata in split_plugin.split(entry=entry): + self._process_entry( + plugins=plugins, dry_run=dry_run, entry=entry, entry_metadata=entry_metadata + ) + def download(self, dry_run: bool = False) -> FileHandlerTransactionLog: """ Performs the subscription download @@ -284,29 +327,14 @@ class Subscription: if isinstance(entry, tuple): entry, entry_metadata = entry - # First, modify the entry with all plugins - for plugin in plugins: - # Break out of this plugin loop if entry is None, it is indicated to not DL it - if (entry := plugin.modify_entry(entry)) is None: - break - - # If entry is None from the broken out loop, continue over the other entries - if entry is None: - continue - - # Then, post-process the entry with all plugins - for plugin in plugins: - optional_plugin_entry_metadata = plugin.post_process_entry(entry) - if optional_plugin_entry_metadata: - entry_metadata.extend(optional_plugin_entry_metadata) - - self._move_entry_files_to_output_directory( - dry_run=dry_run, entry=entry, entry_metadata=entry_metadata - ) - - # Re-save the download archive after each entry is moved to the output directory - if self.maintain_download_archive: - self._enhanced_download_archive.save_download_mappings() + if split_plugin := _get_split_plugin(plugins): + self._process_split_entry( + split_plugin=split_plugin, plugins=plugins, dry_run=dry_run, entry=entry + ) + else: + self._process_entry( + plugins=plugins, dry_run=dry_run, entry=entry, entry_metadata=entry_metadata + ) downloader.post_download(overrides=self.overrides) for plugin in plugins: diff --git a/src/ytdl_sub/utils/chapters.py b/src/ytdl_sub/utils/chapters.py index 53b49700..cb99ae17 100644 --- a/src/ytdl_sub/utils/chapters.py +++ b/src/ytdl_sub/utils/chapters.py @@ -8,7 +8,6 @@ from typing import Optional from typing import Tuple from ytdl_sub.utils.exceptions import ValidationException -from ytdl_sub.utils.ffmpeg import FFMPEG from ytdl_sub.utils.file_handler import FileMetadata @@ -214,17 +213,27 @@ class Chapters: @classmethod def from_embedded_chapters(cls, file_path: str) -> "Chapters": - with BytesIO() as bytes_io: - subprocess.run([ - "-loglevel", "quiet", "-print_format", "json", "-show_chapters", "--", file_path - ], check=True, stdout=bytes_io) - - embedded_chapters = json.load(bytes_io) + proc = subprocess.run( + [ + "ffprobe", + "-loglevel", + "quiet", + "-print_format", + "json", + "-show_chapters", + "--", + file_path, + ], + check=True, + stdout=subprocess.PIPE, + encoding="utf-8", + ) + embedded_chapters = json.loads(proc.stdout) timestamps: List[Timestamp] = [] titles: List[str] = [] - for chapter in embedded_chapters['chapters']: - timestamps.append(Timestamp.from_seconds(int(chapter['start_time']))) - titles.append(chapter['tags']['title']) + for chapter in embedded_chapters["chapters"]: + timestamps.append(Timestamp.from_seconds(int(chapter["start_time"]))) + titles.append(chapter["tags"]["title"]) return Chapters(timestamps=timestamps, titles=titles) diff --git a/tests/e2e/plugins/test_split_by_chapters.py b/tests/e2e/plugins/test_split_by_chapters.py new file mode 100644 index 00000000..deb3b8cf --- /dev/null +++ b/tests/e2e/plugins/test_split_by_chapters.py @@ -0,0 +1,48 @@ +import pytest +from e2e.expected_download import assert_expected_downloads +from e2e.expected_transaction_log import assert_transaction_log_matches + +from ytdl_sub.subscriptions.subscription import Subscription + + +@pytest.fixture +def yt_album_as_chapters_preset_dict(output_directory): + return { + "preset": "yt_album_as_chapters", + "youtube": {"video_url": "youtube.com/watch?v=wtg7AetxuWo"}, + # 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[ext=mp4]", + "postprocessor_args": {"ffmpeg": ["-bitexact"]}, # Must add this for reproducibility + }, + } + + +class TestSplitByChapters: + @pytest.mark.parametrize("dry_run", [True, False]) + def test_video_with_chapters( + self, + youtube_audio_config, + yt_album_as_chapters_preset_dict, + output_directory, + dry_run, + ): + subscription = Subscription.from_dict( + config=youtube_audio_config, + preset_name="split_by_chapters_video", + preset_dict=yt_album_as_chapters_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/split_by_chapters_video.txt", + ) + assert_expected_downloads( + output_directory=output_directory, + dry_run=dry_run, + expected_download_summary_file_name="plugins/split_by_chapters_video.json", + )