ytdl-sub/src/ytdl_sub/downloaders/youtube/split_video.py

202 lines
7.1 KiB
Python

import copy
import re
from pathlib import Path
from typing import Dict
from typing import List
from typing import Tuple
from ytdl_sub.downloaders.youtube.abc import YoutubeDownloader
from ytdl_sub.downloaders.youtube.video 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.ffmpeg import FFMPEG
# 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
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.validators import StringValidator
_SPLIT_TIMESTAMP_REGEX = re.compile(r"^((?:\d\d:)?(?:\d:)?(?:\d)?\d:\d\d) (.+)$")
def _split_video_uid(source_uid: str, idx: int) -> str:
return f"{source_uid}___{idx}"
def _split_video_ffmpeg_cmd(
input_file: str, output_file: str, timestamps: List[Timestamp], idx: int
) -> List[str]:
timestamp_begin = timestamps[idx].standardized_str
timestamp_end = timestamps[idx + 1].standardized_str if idx + 1 < len(timestamps) else ""
cmd = ["-i", input_file, "-ss", timestamp_begin]
if timestamp_end:
cmd += ["-to", timestamp_end]
cmd += ["-vcodec", "copy", "-acodec", "copy", output_file]
return cmd
class YoutubeSplitVideoDownloaderOptions(YoutubeVideoDownloaderOptions):
r"""
Downloads a single youtube video, then splits in to separate videos using a file containing
timestamps. Each separate video will be formatted as if it was downloaded from a playlist.
This download strategy is intended for CLI usage performing a one-time download of a video,
not a subscription.
Usage:
.. code-block:: yaml
presets:
example_preset:
youtube:
# required
download_strategy: "split_video"
video_url: "youtube.com/watch?v=VMAPTo7RVDo"
split_timestamps: path/to/timestamps.txt
CLI usage:
.. code-block:: bash
ytdl-sub dl \
--preset "example_preset" \
--youtube.video_url "youtube.com/watch?v=VMAPTo7RVDo" \
--youtube.split_timestamps "path/to/timestamps.txt"
``split_timestamps`` file format:
.. code-block:: markdown
0:00 Intro
0:24 Blackwater Park
10:23 Bleak
16:39 Jokes
1:02:23 Ending
The above will create 5 videos in total. The last timestamp, in this example, would create a
video starting at ``1:02:23`` and end at Youtube video's ending.
"""
_required_keys = {"video_url", "split_timestamps"}
def __init__(self, name, value):
super().__init__(name, value)
self._split_timestamps = self._validate_key("split_timestamps", StringValidator).value
@property
def split_timestamps(self) -> str:
"""
Required. The path to the file containing the split timestamps.
"""
return self._split_timestamps
class YoutubeSplitVideoDownloader(
YoutubeDownloader[YoutubeSplitVideoDownloaderOptions, YoutubePlaylistVideo]
):
downloader_options_type = YoutubeSplitVideoDownloaderOptions
downloader_entry_type = YoutubePlaylistVideo
supports_download_archive = False
supports_subtitles = False
@classmethod
def ytdl_option_defaults(cls) -> Dict:
"""
Default `ytdl_options`_ for ``split_video``
.. code-block:: yaml
ytdl_options:
ignoreerrors: True # ignore errors like hidden videos, age restriction, etc
"""
return dict(
super().ytdl_option_defaults(),
**{"break_on_existing": True},
)
def _create_split_video_entry(
self, source_entry_dict: Dict, title: str, idx: int, chapters: Chapters
) -> Tuple[YoutubePlaylistVideo, FileMetadata]:
"""
Runs ffmpeg to create the split video
"""
entry_dict = copy.deepcopy(source_entry_dict)
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)
# 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"]
timestamp_begin = chapters.timestamps[idx].readable_str
timestamp_end = Timestamp(source_entry_dict["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,
)
def download(self) -> List[Tuple[YoutubePlaylistVideo, FileMetadata]]:
"""Download a single Youtube video, then split it into multiple videos"""
split_videos_and_metadata: List[Tuple[YoutubePlaylistVideo, FileMetadata]] = []
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)
# convert the entry thumbnail early so we do not have to guess the thumbnail extension
# when copying it
if not self.is_dry_run:
convert_download_thumbnail(entry=entry)
for idx, title in enumerate(chapters.titles):
new_uid = _split_video_uid(source_uid=entry.uid, idx=idx)
if not self.is_dry_run:
# Get the input/output file paths
input_file = entry.get_download_file_path()
output_file = str(Path(self.working_directory) / f"{new_uid}.{entry.ext}")
# Run ffmpeg to create the split the video
FFMPEG.run(
_split_video_ffmpeg_cmd(
input_file=input_file,
output_file=output_file,
timestamps=chapters.timestamps,
idx=idx,
)
)
# Copy the original vid thumbnail to the working directory with the new uid. This so
# downstream logic thinks this split video has its own thumbnail
FileHandler.copy(
src_file_path=entry.get_download_thumbnail_path(),
dst_file_path=Path(self.working_directory) / f"{new_uid}.{entry.thumbnail_ext}",
)
# 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
)
)
return split_videos_and_metadata