Add a `--split-chapters` option to `youtube-dl`

ytdl_chapters.py

#!/usr/bin/env python
#-*- coding: utf-8 -*-

# hacky way of adding a new option to the parser, and subsequently working out what the filename is called

import subprocess
import sys
import os
import glob
import re
import optparse
import json

from typing import List

import youtube_dl
import youtube_dl.options

filenames = []
split_enabled = False
retcode = 0

__old_YoutubeDL = youtube_dl.YoutubeDL


class new_YoutubeDL(__old_YoutubeDL):
    def prepare_filename(self, *args, **kwargs):
        global filenames
        filename = super().prepare_filename(*args, **kwargs)
        filenames.append(filename)
        return filename


__old_OptionGroup = optparse.OptionGroup


class new_OptionGroup(__old_OptionGroup):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if "Download Options" in args:
            self.add_option(
                "--split-chapters",
                dest="splitchapters",
                action="store_true",
                default=False,
                help="Splits download into individual videos defined by YouTube chapters (implies --write-info-json).",
            )


__old_parseOpts = youtube_dl.options.parseOpts


def new_parseOpts(*args, **kwargs):
    global split_enabled
    parser, opts, args = __old_parseOpts(*args, **kwargs)
    if opts.splitchapters:
        opts.writeinfojson = True
        split_enabled = True

    return parser, opts, args


def split_by_chapters(filenames: List[str]) -> None:
    print("Splitting download by chapters:")

    for file in filenames:
        # check files exist
        media = "".join(file.split(".")[:-1])

        files = glob.glob(f"./{media}.*")

        try:
            json_file = [i for i in files if ".json" in i][0]
            media_files = [i for i in files if i != json_file]
        except:
            pass
        
        else:
            if json_file and media_files:
                if len(media_files) > 1:
                    # ask user to resolve ambiguity
                    for i, f in enumerate(media_files):
                        print(f"{i+1} -- {f}")

                    s = 0
                    while not (0 < s <= len(media_files)):
                        s = input(
                            f"Select a number to use corresponding file (1-{len(media_files)}):\n"
                        )
                        try:
                            s = int(s)
                        except:
                            s = 0

                    media_file = media_files[s - 1]
                else:
                    media_file = media_files[0]

                ext = media_file.split(".")[-1]
                break
    else:
        print(
            "ERROR: Bad file globbing for filenames."
        )  # this error message makes no sense to anyone but maybe me
        return

    if re.search(r"^(mp3|webm|mp4)$", ext):
        pass
    else:
        print("ERROR: No trimming function known for this file format.")
        return

    with open(json_file, "r") as f:
        content = f.read()

    try:
        os.mkdir(media)
    except FileExistsError:
        print("WARNING: Output directory already exists.")

    chapters = json.loads(content)["chapters"]
    for i, c in enumerate(chapters):
        ss = c["start_time"]
        dt = c["end_time"] - ss
        title = c["title"]

        outfile = f"{media}/{i+1}_{title}.{ext}"

        cmd = [
            "ffmpeg",
            "-i",
            f"\"{media_file}\"",
            "-ss",
            str(ss),
            "-t",
            str(dt),
            f"\"{outfile}\"",
        ]
        cmd = " ".join(cmd)
        
        if os.path.isfile(outfile):
            print(f"File {outfile} already exists... skipping.")
            continue

        subprocess.run(cmd, shell=True, check=True)

    print(f"Trimmed and created {len(chapters)} chapters in '{media}'.")


if __name__ == "__main__":

    __old_exit = sys.exit

    def new_exit(*args, **kwargs):
        global retcode
        retcode = args[0]

    # patch
    sys.modules["sys"].exit = new_exit
    sys.modules["youtube_dl"].YoutubeDL = new_YoutubeDL
    sys.modules["youtube_dl"].parseOpts = new_parseOpts
    sys.modules["youtube_dl.options"].optparse.OptionGroup = new_OptionGroup

    # do the download
    youtube_dl.main(sys.argv[1:])

    if split_enabled:
        split_by_chapters(filenames)

    __old_exit(retcode)