156 lines
4.7 KiB
Python
156 lines
4.7 KiB
Python
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.file_handler import FileHandler
|
|
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
|
|
def _ensure_installed(cls):
|
|
try:
|
|
subprocess.check_output(["which", "ffmpeg"])
|
|
except subprocess.CalledProcessError as subprocess_error:
|
|
raise ValidationException(
|
|
"Trying to use a feature which requires ffmpeg, but it cannot be found"
|
|
) from subprocess_error
|
|
|
|
@classmethod
|
|
def run(cls, ffmpeg_args: List[str]) -> None:
|
|
"""
|
|
Runs an ffmpeg command. Should not include 'ffmpeg' as the beginning argument.
|
|
|
|
Parameters
|
|
----------
|
|
ffmpeg_args:
|
|
Arguments to pass to ffmpeg. Each one will be separated by a space.
|
|
"""
|
|
cls._ensure_installed()
|
|
|
|
cmd = ["ffmpeg"]
|
|
cmd.extend(ffmpeg_args)
|
|
logger.debug("Running %s", " ".join(cmd))
|
|
with Logger.handle_external_logs(name="ffmpeg"):
|
|
subprocess.run(cmd, check=True, capture_output=True)
|
|
|
|
|
|
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 set_ffmpeg_metadata_chapters(
|
|
file_path: str, chapters: Optional[Chapters], file_duration_sec: int
|
|
) -> None:
|
|
"""
|
|
Sets ffmetadata chapters to a file. Note that this will (I think) wipe all prior
|
|
metadata.
|
|
|
|
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 chapters:
|
|
lines += _create_metadata_chapters(chapters=chapters, file_duration_sec=file_duration_sec)
|
|
|
|
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()
|
|
|
|
FFMPEG.run(
|
|
[
|
|
"-i",
|
|
file_path,
|
|
"-i",
|
|
metadata_file.name,
|
|
"-map",
|
|
"0",
|
|
"-map_chapters",
|
|
"1",
|
|
"-bitexact", # for reproducibility
|
|
"-codec",
|
|
"copy",
|
|
output_file_path,
|
|
]
|
|
)
|
|
FileHandler.move(output_file_path, file_path)
|
|
|
|
|
|
def add_ffmpeg_metadata_key_values(file_path: str, key_values: Dict[str, str]) -> None:
|
|
"""
|
|
Parameters
|
|
----------
|
|
file_path
|
|
File to add metadata key/values to
|
|
key_values
|
|
The key/values to add
|
|
"""
|
|
file_path_ext = file_path.split(".")[-1]
|
|
output_file_path = f"{file_path}.out.{file_path_ext}"
|
|
|
|
ffmpeg_args = ["-i", file_path, "-map", "0"]
|
|
for key, value in key_values.items():
|
|
ffmpeg_args.extend(["-metadata", f"{key}={value}"])
|
|
ffmpeg_args.extend(["-codec", "copy", output_file_path])
|
|
|
|
FFMPEG.run(ffmpeg_args)
|
|
FileHandler.move(output_file_path, file_path)
|