script-dump/rip-py/rip_py/utils.py

import math
import os
from pathlib import Path
import shutil
import subprocess
import tempfile


def with_fifo(func):
    def f(*args, **kwargs):
        with tempfile.TemporaryDirectory() as tmpdir:
            fifo_path = str(Path(tmpdir) / 'rip.fifo')
            try:
                os.mkfifo(fifo_path)
            except FileExistsError:
                pass
            kwargs['fifo_path'] = fifo_path
            return func(*args, **kwargs)
    return f


def convert_chapters_to_mkvmerge_format(chapters):
    """Convert a list of chapter timestamps into a string (of multiple lines) as accepted by mkvmerge"""
    # determine how many zeroes we have as a prefix
    number_of_zeros = 0
    if len(chapters) > 1:
        number_of_zeros = int(math.log10(len(chapters) - 1))

    chapter_template = 'CHAPTER{{:0{}d}}'.format(number_of_zeros)

    lines = []
    for i, timestamp in enumerate(chapters):
        chapter_str = chapter_template.format(i)

        # if we start with e.g. 0:21, we need another 0 to be compatible
        if len(timestamp.split(':')[0]) == 1:
            timestamp = '0' + timestamp

        # if we have a higher resolution than 1 millisecond, we need to cut
        # some chars
        if len(timestamp.split('.')[-1]) > 3:
            strip_length = len(timestamp.split('.')[-1]) - 3
            timestamp = timestamp[:-strip_length]

        lines.append(f"{chapter_str}={timestamp}")
        lines.append(f"{chapter_str}NAME=")

    return os.linesep.join(lines)


def set_chapters(mkv_file, chapters):
    """Set the chapters in the given MKV file

    `chapters` is a list of chapter start times.
    """
    if not Path(mkv_file).exists():
        raise RuntimeError('MKV file does not exist.')

    if not shutil.which('mkvpropedit'):
        raise RuntimeError('Cannot find "mkvpropedit". Please install the "mkvtoolnix" package.')

    with tempfile.NamedTemporaryFile() as f:
        mkvmerge_chapters = convert_chapters_to_mkvmerge_format(chapters)
        f.write(mkvmerge_chapters.encode('utf-8'))
        f.flush()

        mkvpropedit_cmd = ['mkvpropedit', mkv_file, '--chapters', f.name]
        subprocess.check_call(mkvpropedit_cmd)


def get_parameterized_ffmpeg_cmd(media_type, quality, extra_opts, use_gpu,
                                 audio_map, subtitle_map, pre_read=1024,
                                 video_filters=None, scale=None):
    """Build an ffmpeg command based on the given parameters

    Returns an function taking input file and output file as arguments to build
    the final command.
    """
    if scale is not None:
        if video_filters is None:
            video_filters = []
        video_filters.append(f"scale={scale}")

    ffmpeg_cmd = []
    ffmpeg_cmd.append(['ffmpeg',
                       '-probesize', '{}M'.format(pre_read),
                       '-analyzeduration', '{}M'.format(pre_read),
                       '-hwaccel', 'auto',
                       '-i'])

    ffmpeg_cmd.append(extra_opts)
    ffmpeg_cmd[-1].extend(['-map', '0:v:0'])
    if '-c:a' not in ffmpeg_cmd[-1]:
        ffmpeg_cmd[-1].extend(['-c:a', 'copy'])
    ffmpeg_cmd[-1].extend(['-c:s', 'copy'])
    if use_gpu:
        ffmpeg_cmd[-1] += [
            '-c:v', 'h264_nvenc',
            '-rc', 'vbr_hq',
            '-rc-lookahead', '20',
            '-cq', '{}'.format(quality),
            # '-maxrate:v', '3M',
            # '-b:v', '8M', '-maxrate:v', '10M',
            '-preset:v', 'slow']
    else:
        ffmpeg_cmd[-1] += [
            '-c:v', 'h264',
            '-crf', '{}'.format(quality),
            '-preset:v', 'slower',
            '-x264-params', 'opencl=true',
            # '-x264-params', 'keyint=240:min-keyint=20',
        ]

    ffmpeg_cmd[-1] += [
        '-pix_fmt', 'yuv420p',
        '-movflags', '+faststart',
        # '-profile:v', 'baseline',
        # '-level', '3.0',
    ]

    if video_filters:
        ffmpeg_cmd[-1] += [
            '-vf', ','.join(video_filters)
        ]

    if audio_map:
        for i, (channel, lang) in enumerate(audio_map):
            if channel < 0x80:
                map_target = '0:a:{}'.format(channel)
            else:
                map_target = 'i:0x{:x}'.format(channel)
            ffmpeg_cmd[-1].extend([
                '-map', map_target,
                '-metadata:s:a:{}'.format(i), 'language={}'.format(lang)
            ])
    else:
        ffmpeg_cmd[-1].extend(['-map', '0:a:0', '-metadata:s:a:0', 'language=de'])

    if subtitle_map:
        for i, (channel, lang) in enumerate(subtitle_map):
            if channel < 0x20:
                map_target = '0:s:{}'.format(channel)
            else:
                map_target = 'i:0x{:x}'.format(channel)
            ffmpeg_cmd[-1].extend([
                '-map', map_target,
                '-metadata:s:s:{}'.format(i), 'language={}'.format(lang)
            ])

    # TODO we are probably able to do forced subtitles. we probably want that

    # we can set a default stream. make this available somehow
    #  -disposition:a:1 default

    if media_type and not use_gpu:
        ffmpeg_cmd[-1].extend(['-tune:v', media_type])

    def parameterized(output_file, input_file):
        return ffmpeg_cmd[0] + [input_file] + ffmpeg_cmd[1] + [output_file]

    return parameterized


def get_ffmpeg_cmd(output_file, input_path, media_type, quality, extra_opts,
                   use_gpu, audio_map, subtitle_map, video_filters=None,
                   scale=None):
    cmd = get_parameterized_ffmpeg_cmd(media_type, quality, extra_opts,
                                       use_gpu, audio_map, subtitle_map,
                                       video_filters=video_filters, scale=scale)
    cmd = cmd(output_file=output_file, input_file=input_path)
    print(cmd)
    return cmd


@with_fifo
def rip_title(title, output_dir, mpv_cmd, ffmpeg_cmd, fifo_path,
              output_file_suffix='.mkv'):
    filename = f"{title}{output_file_suffix}"
    out_file = output_dir / filename

    mpv_cmd = mpv_cmd(title, fifo_path)
    # we need to redirect stdout, because otherwise we endup botching the terminal input somehow
    mpv_proc = subprocess.Popen(mpv_cmd, stdout=subprocess.DEVNULL)

    ffmpeg_cmd = ffmpeg_cmd(output_file=out_file, input_file=fifo_path)
    print(ffmpeg_cmd)
    subprocess.check_call(ffmpeg_cmd)

    # mpv should exit if ffmpeg is gone, because the FIFO got closed
    mpv_proc.communicate()

    return out_file