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

from datetime import datetime
import functools
import itertools
from pathlib import Path
import shlex
import subprocess

import click

from rip_py import _bluray, _dvd, utils


FILE_TYPE = click.Path(file_okay=True, dir_okay=False, exists=True, path_type=Path)
NEW_FILE_TYPE = click.Path(file_okay=True, dir_okay=False, exists=False, path_type=Path)
DIRECTORY_TYPE = click.Path(file_okay=False, dir_okay=True, exists=True, path_type=Path)
DEVICE_OR_DIRECTORY_TYPE = click.Path(file_okay=True, dir_okay=True, exists=True, path_type=Path)


def scale_convert(ctx, param, value):
    if value == '720p':
        return '-1:720'
    if value == '480p':
        return '-1:480'
    return value


def _rip(module, output_dir, titles, device, media_type, keep_audio, quality,
         start_at_title, end_at_title, ffmpeg_opts, deinterlace, audio_map,
         subtitle_map, pre_read_size, scale=None):
    if end_at_title >= 0 and start_at_title > end_at_title:
        raise click.ClickException('--start-at-title cannot be larger than --end-at-title.')

    ffmpeg_opts = shlex.split(ffmpeg_opts) if ffmpeg_opts is not None else []

    video_filters = []
    if deinterlace:
        video_filters.append('yadif')

    if titles:
        titles = sorted(t for t in titles if t >= start_at_title)
    else:
        titles = module.get_num_titles(device)
        if end_at_title >= 0:
            if end_at_title > titles:
                raise click.ClickException(f"--end-at-title {end_at_title} larger than number of titles {titles}")
            # we'll put it into a Python range(), so we need to add 1 to make
            # it include that title
            end_at_title += 1
        else:
            end_at_title = titles
        titles = range(start_at_title, end_at_title)

    for title in titles:
        if keep_audio:
            audio_map = module.get_title_audio_map(device, title)

        use_gpu = False
        ffmpeg_cmd = utils.get_parameterized_ffmpeg_cmd(
            media_type, quality, ffmpeg_opts, use_gpu, audio_map, subtitle_map,
            pre_read=pre_read_size, video_filters=video_filters, scale=scale)
        mpv_cmd = module.get_parameterized_mpv_dump_cmd(device)

        out_file = utils.rip_title(title, output_dir, mpv_cmd, ffmpeg_cmd)

        chapters = module.get_title_chapters(device, title)
        if chapters:
            utils.set_chapters(out_file, chapters)


@click.group()
def cli():
    pass


@cli.group()
@click.option('-d', '--device', type=DEVICE_OR_DIRECTORY_TYPE,
              default='/dev/dvd', show_default=True)
@click.pass_context
def dvd(ctx, device):
    ctx.ensure_object(dict)
    ctx.obj['device'] = str(device)


def pass_device(fn):
    @functools.wraps(fn)
    @click.pass_obj
    def wrapped(obj, *args, **kwargs):
        kwargs['device'] = obj['device']
        return fn(*args, **kwargs)

    return wrapped


@dvd.group('info')
def dvd_info():
    ...


@dvd_info.command('audios')
@click.option('-t', '--title', type=int, required=True)
@pass_device
def dvd_info_audios(device, title):
    """Show the audio streams and their languages for a title"""
    for i, (channel, lang) in enumerate(_dvd.get_title_audio_map(device, title)):
        print(f"{i}: {lang} aid probably {channel} (0x{channel:x})")


@dvd_info.command('chapters')
@click.option('-t', '--title', type=int, required=True)
@pass_device
def dvd_info_chapters(title, device):
    """Show the chapters for a title"""
    print(','.join(_dvd.get_title_chapters(device, title)))


@dvd_info.command('ffprobe')
@click.option('-t', '--title', type=int, required=True)
@pass_device
def dvd_info_ffprobe(title, device):
    """Run ffprobe on the given title

    Output is saved into a file named like the title number in the given directory.
    """
    ffmpeg_cmd = [['ffprobe', '-i']]

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

    mpv_cmd = _dvd.get_parameterized_mpv_dump_cmd(device)
    utils.rip_title(title, Path('/tmp'), mpv_cmd)


@dvd_info.command('subtitles')
@click.option('-t', '--title', type=int, required=True)
@pass_device
def dvd_info_subtitles(device, title):
    """Show the subtitle streams and their languages for a title"""
    for i, (channel, lang) in enumerate(_dvd.get_title_subtitle_map(device, title)):
        print(f"{i}: {lang} with id {channel} (0x{channel:x})")


@dvd_info.command('titles')
@pass_device
def dvd_info_titles(device):
    """Show the titles and their length"""
    for i, t in enumerate(_dvd.get_titles(device)):
        print(f"{i}: {t}")


@dvd.command()
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('-t', '--title', type=int, required=True)
@click.option('--pre-read-size', type=int, default=1024, show_default=True)
@pass_device
def get_subtitles_file(output_dir, title, device, pre_read_size):
    """Write out the subtitles of the given title into the given directory"""
    # ffmpeg_cmd = ffmpeg_cmd(output_file=out_file, input_file=fifo_path)
    # FIXME can we move this out of the `cli` module?
    ffmpeg_cmd = []
    ffmpeg_cmd.append(['ffmpeg',
                       '-probesize', '{}M'.format(pre_read_size),
                       '-analyzeduration', '{}M'.format(pre_read_size),
                       '-hwaccel', 'auto',
                       '-i'])
    ffmpeg_cmd.append([
        '-map', '0:v:0',
        '-c:v', 'h264',
        '-crf', '75',
        '-preset:v', 'veryfast',
        '-x264-params', 'opencl=true',
        # '-x264-params', 'keyint=240:min-keyint=20',
        '-c:s', 'copy',
    ])
    for i, (channel, lang) in enumerate(_dvd.get_title_subtitle_map(device, title)):
        lang = f"{lang} - 0x{channel:x}"
        ffmpeg_cmd[-1] += [
            '-map', f"i:0x{channel:x}",
            f"-metadata:s:s:{i}", f"language={lang}"
        ]

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

    mpv_cmd = _dvd.get_parameterized_mpv_dump_cmd(device)

    utils.rip_title(title, output_dir, mpv_cmd, parameterized)


@dvd.command()
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('-t', '--title', type=int)
@click.option('--no-dump', is_flag=True)
@click.option('--no-audios', is_flag=True)
@click.option('--no-subtitles', is_flag=True)
@click.option('--no-chapters', is_flag=True)
def dump(title, output_dir, device, no_dump, no_audios, no_subtitles,
         no_chapters):
    """Just dump the stream and additional infos

    It's basically ripping without converting."""
    if title is None:
        titles = range(0, _dvd.get_num_titles(device))
    else:
        titles = [title]

    for title in titles:
        if not no_dump:
            dump_file = str(output_dir / f"{title}.dump")
            mpv_cmd = _dvd.get_mpv_dump_cmd(title, dump_file, device)

            subprocess.check_call(mpv_cmd)

        if not no_audios:
            audios_file = str(output_dir / f"{title}.audios")
            with open(audios_file, 'w') as f:
                audios = [f"{c}: {l}"
                          for c, l in _dvd.get_title_audio_map(device, title)]
                f.write('\n'.join(audios))

        if not no_subtitles:
            subtitles_file = str(output_dir / f"{title}.subs")
            with open(subtitles_file, 'w') as f:
                subtitles = [f"{c}: {l}"
                             for c, l in _dvd.get_title_subtitle_map(device, title)]
                f.write('\n'.join(subtitles))

        if not no_chapters:
            chapters_file = str(output_dir / f"{title}.chapters")
            with open(chapters_file, 'w') as f:
                chapters = _dvd.get_title_chapters(device, title)
                f.write(','.join(chapters))


@dvd.command('rip')
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('--media-type', type=click.Choice(['film', 'animation']))
@click.option('--keep-audio', is_flag=True)
@click.option('--audio-map', multiple=True, type=(int, str),
              help="Keep specific audio. Order is important. Takes first "
                   "audio stream and puts it at stream INT in the new file "
                   "with language STR.")
@click.option('--subtitle-map', multiple=True, type=(int, str),
              help="Keep the specific subtitles. Order is important. Takes "
                   "stream defined as INT into place according to order with "
                   "language STR.")
@click.option('--quality', type=int, default=22, show_default=True)
@click.option('--start-at-title', type=int, default=0, show_default=True)
@click.option('--end-at-title', type=int, default=-1)
@click.option('-t', '--title', 'titles', multiple=True, type=int)
@click.option('--deinterlace', is_flag=True)
@click.option('--pre-read-size', type=int, default=1024, show_default=True)
@click.option('--ffmpeg-opts')
@pass_device
def dvd_rip(output_dir, titles, device, media_type, keep_audio, quality,
            start_at_title, end_at_title, ffmpeg_opts, deinterlace, audio_map,
            subtitle_map, pre_read_size):
    _rip(_dvd, output_dir, titles, device, media_type, keep_audio, quality,
         start_at_title, end_at_title, ffmpeg_opts, deinterlace, audio_map,
         subtitle_map, pre_read_size)


@dvd.command('rip-audio')
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('-t', '--title', type=int, required=True)
@pass_device
def dvd_rip_audio(device, output_dir, title):
    """Rip the given title's audio from DVD"""
    def partial(output_file, input_file):
        return ['ffmpeg',
                '-i', input_file,
                '-map', '0:a:0',
                '-c:a', 'flac',
                output_file]

    mpv_cmd = _dvd.get_parameterized_mpv_dump_cmd(device)

    utils.rip_title(title, output_dir, mpv_cmd, partial,
                    output_file_suffix='.flac')


@dvd.command('find-right-title')
@click.option('--start-title', type=int, default=0, show_default=True)
@click.option('--end-title', type=int)
@pass_device
def dvd_find_right_title(start_title, end_title, device):
    """Try to find a title without duplicate chapter lengths"""
    if end_title and start_title > end_title:
        msg = '--start-title must not be bigger than --end-title'
        raise click.ClickException(msg)

    titles = _dvd.get_num_titles(device)
    if start_title > titles:
        raise RuntimeError('DVD has less titles than start title.')

    if end_title:
        end_title = min(titles - 1, end_title)
    else:
        end_title = titles - 1

    title_chapter_lengths = {}
    for title in range(start_title, end_title + 1):
        chapters = _dvd.get_title_chapters(device, title)

        timestamps = []
        for timestamp in chapters:
            timestamp, microseconds = timestamp.split('.')
            dt = datetime.strptime(timestamp, '%H:%M:%S')
            dt = dt.replace(microsecond=int(microseconds) * 1000)
            timestamps.append(dt)

        lengths = []
        if len(timestamps) > 1:
            for i, timestamp in enumerate(timestamps[1:]):
                lengths.append(timestamp - timestamps[i])

        if not lengths:
            continue

        title_chapter_lengths[title] = lengths

    # remove titles with duplicate chapters by looking at the length
    title_chapter_lengths = {
        title: lengths for title, lengths in title_chapter_lengths.items()
        if len(lengths) == len(set(lengths))}

    if title_chapter_lengths:
        print('The titles {} contain no duplicate chapter lengths.'.format(
              ', '.join(str(i) for i in title_chapter_lengths.keys())))
    else:
        print('No title without duplicates found.')


@cli.group()
@click.option('-d', '--device', type=DEVICE_OR_DIRECTORY_TYPE,
              default='/dev/dvd', show_default=True)
@click.pass_context
def bluray(ctx, device):
    ctx.ensure_object(dict)
    ctx.obj['device'] = str(device)


@bluray.group('info')
def bluray_info():
    ...


@bluray_info.command('audios')
@click.option('-t', '--title', type=int, required=True)
@pass_device
def bluray_info_audios(device, title):
    """Show the audio streams and their languages for a title"""
    for i, (channel, lang) in enumerate(_bluray.get_title_audio_map(device, title)):
        print(f"{i}: {lang} aid probably {channel} (0x{channel:x})")


@bluray_info.command('chapters')
@click.option('-t', '--title', type=int, required=True)
@pass_device
def bluray_info_chapters(title, device):
    """Show the chapters for a title"""
    print(','.join(_bluray.get_title_chapters(device, title)))


@bluray_info.command('ffprobe')
@click.option('-t', '--title', type=int, required=True)
@pass_device
def bluray_info_ffprobe(title, device):
    """Run ffprobe on the given title

    Output is saved into a file named like the title number in the given directory.
    """
    # ffmpeg_cmd = [['ffprobe', '-probesize', '1024M', '-analyzeduration', '1024M', '-i']]
    ffmpeg_cmd = [['ffprobe', '-i']]

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

    mpv_cmd = _bluray.get_parameterized_mpv_dump_cmd(device)

    utils.rip_title(title, Path('/tmp'), mpv_cmd, parameterized)


@bluray_info.command('main-title')
@pass_device
def bluray_info_main_title(device):
    """Show the main title as advertised by the disk"""
    title, playlist = _bluray.get_main_title(device)
    print(f"Main title: {title} with playlist {playlist}")


@bluray_info.command('subtitles')
@click.option('-t', '--title', type=int, required=True)
@pass_device
def bluray_info_subtitles(device, title):
    """Show the subtitle streams and their languages for a title"""
    for i, (channel, lang) in enumerate(_bluray.get_title_subtitle_map(device, title)):
        print(f"{i}: {lang} with id {channel} (0x{channel:x})")


@bluray_info.command('titles')
@pass_device
def bluray_info_titles(device):
    """Show the titles and their length"""
    for i, t in enumerate(_bluray.get_titles(device)):
        print(f"{i}: {t}")


@bluray.command('rip')
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('--media-type', type=click.Choice(['film', 'animation']))
@click.option('--keep-audio', is_flag=True)
@click.option('--audio-map', multiple=True, type=(int, str),
              help="Keep specific audio. Order is important. Takes first "
                   "audio stream and puts it at stream INT in the new file "
                   "with language STR.")
@click.option('--subtitle-map', multiple=True, type=(int, str),
              help="Keep the specific subtitles. Order is important. Takes "
                   "stream defined as INT into place according to order with "
                   "language STR.")
@click.option('--quality', type=int, default=22, show_default=True)
@click.option('--start-at-title', type=int, default=0, show_default=True)
@click.option('--end-at-title', type=int, default=-1)
@click.option('-t', '--title', 'titles', multiple=True, type=int)
@click.option('--deinterlace', is_flag=True)
@click.option('--pre-read-size', type=int, default=1024, show_default=True)
@click.option('--ffmpeg-opts')
@click.option('--scale', type=click.Choice(['720p', '480p']), callback=scale_convert,
              help='Re-scale the video to the given size')
@pass_device
def bluray_rip(output_dir, titles, device, media_type, keep_audio, quality,
               start_at_title, end_at_title, ffmpeg_opts, deinterlace,
               audio_map, subtitle_map, pre_read_size, scale):
    _rip(_bluray, output_dir, titles, device, media_type, keep_audio, quality,
         start_at_title, end_at_title, ffmpeg_opts, deinterlace, audio_map,
         subtitle_map, pre_read_size, scale=scale)


@bluray.command('rip-audio')
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('-t', '--title', type=int, required=True)
@pass_device
def bluray_rip_audio(device, output_dir, title):
    """Rip the given title's audio from DVD"""
    def partial(output_file, input_file):
        return ['ffmpeg',
                '-i', input_file,
                '-map', '0:a:0',
                '-c:a', 'flac',
                output_file]

    mpv_cmd = _bluray.get_parameterized_mpv_dump_cmd(device)

    utils.rip_title(title, output_dir, mpv_cmd, partial,
                    output_file_suffix='.flac')


@cli.group()
def video():
    ...


@video.command('convert')
@click.option('-i', '--input-file', type=FILE_TYPE, required=True)
@click.option('-o', '--output-file', type=NEW_FILE_TYPE, required=True)
@click.option('--media-type', type=click.Choice(['film', 'animation']))
@click.option('--quality', type=int, default=18, show_default=True)
@click.option('--deinterlace', is_flag=True)
@click.option('--ffmpeg-opts')
@click.option('--scale', type=click.Choice(['720p', '480p']), callback=scale_convert,
              help='Re-scale the video to the given size')
@click.option('--audio-map', multiple=True, type=(int, str),
              help="Keep specific audio. Order is important. Takes first "
                   "audio stream and puts it at stream INT in the new file "
                   "with language STR.")
@click.option('--subtitle-map', multiple=True, type=(int, str),
              help="Keep the specific subtitles. Order is important. Takes "
                   "stream defined as INT into place according to order with "
                   "language STR.")
@click.option('-c', '--chapters')
def video_convert(input_file, output_file, media_type, quality, deinterlace,
                  ffmpeg_opts, audio_map, subtitle_map, chapters, scale):
    ffmpeg_opts = shlex.split(ffmpeg_opts) if ffmpeg_opts is None else []

    video_filters = []
    if deinterlace:
        video_filters.append('yadif')

    if chapters:
        chapters = chapters.split(',')

    ffmpeg_cmd = utils.get_ffmpeg_cmd(output_file, input_file, media_type,
                                      quality, ffmpeg_opts, False, audio_map,
                                      subtitle_map, video_filters=video_filters,
                                      scale=scale)

    subprocess.check_call(ffmpeg_cmd)

    if chapters:
        utils.set_chapters(output_file, chapters)


@cli.command('set-chapters')
@click.option('-o', '--output-file', type=FILE_TYPE, required=True)
@click.option('-c', '--chapters', type=FILE_TYPE, required=True,
              help='File containing a comma-separated list of timestamps')
def set_chapters(output_file, chapters):
    """Set chapters in given file according to the timestamps in the chapters file"""
    chapters = chapters.split(',')

    utils.set_chapters(output_file, chapters)


@video.command('split-file')
@click.option('-c', '--chapters-file', type=FILE_TYPE, required=True,
              help='File with timestamps separated by comma')
@click.option('-i', '--input-file', type=FILE_TYPE, required=True)
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('-n', '--number', type=int,
              help='Split out only one chapter given by this number.')
@click.option('--reencode', is_flag=True,
              help='Don\'t use -c:v copy and -c:a copy')
@click.option('--media-type', type=click.Choice(['film', 'animation']),
              help='Only relevant with --reencode')
@click.option('--quality', type=int, default=18, show_default=True,
              help='Only relevant with --reencode')
@click.option('--ffmpeg-opts',
              help='Only relevant with --reencode')
@click.option('--scale', type=click.Choice(['720p', '480p']), callback=scale_convert,
              help='Re-scale the video to the given size')
@click.option('--audio-map', multiple=True, type=(int, str),
              help="Keep specific audio. Only relevant with --reencode."
                   "Order is important. Takes first "
                   "audio stream and puts it at stream INT in the new file "
                   "with language STR.")
@click.option('--subtitle-map', multiple=True, type=(int, str),
              help="Keep the specific subtitles. Order is important. Takes "
                   "stream defined as INT into place according to order with "
                   "language STR.")
def utils_split_file(output_dir, input_file, chapters_file, number, reencode,
                     media_type, quality, ffmpeg_opts, audio_map, subtitle_map,
                     scale):
    if ffmpeg_opts is None:
        ffmpeg_opts = []
    else:
        ffmpeg_opts = shlex.split(ffmpeg_opts)

    extension = input_file.split('.')[-1]

    with open(chapters_file) as f:
        chapters = f.read().strip().split(',')

    for i, (start, end) in enumerate(itertools.zip_longest(chapters, chapters[1:]), start=1):
        if number is not None and i != number:
            continue

        extra_opts = ffmpeg_opts + ['-ss', start]
        if end is not None:
            extra_opts.extend(['-to', end])

        out_path = output_dir / f"{i}.{extension}"

        if not reencode:
            ffmpeg_cmd = ['ffmpeg'] + extra_opts + [
                '-i', input_file,
                '-c:a', 'copy',
                '-c:v', 'copy',
                '-movflags', '+faststart',
                str(out_path)
            ]
        else:
            use_gpu = False
            ffmpeg_cmd = utils.get_ffmpeg_cmd(out_path, input_file, media_type,
                                              quality, extra_opts, use_gpu,
                                              audio_map, subtitle_map,
                                              scale=scale)
        subprocess.check_call(ffmpeg_cmd)


@cli.group()
def audio():
    ...


@audio.command('split-file')
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('-c', '--chapters-file', type=FILE_TYPE, required=True,
              help='File with timestamps separated by comma')
@click.option('-i', '--input-file', type=FILE_TYPE, required=True)
def audio_split_file(output_dir, chapters_file, input_file):
    """Split an audio file into different flac files based on timestamps in the chapters file"""
    with open(chapters_file) as f:
        chapters = f.read().strip().split(',')

    times_ten = 0
    while len(chapters) // (10 ** times_ten) > 9:
        times_ten += 1

    name_template = '{:0' + str(times_ten + 1) + 'd}.flac'
    for i, (start, end) in enumerate(itertools.zip_longest(chapters, chapters[1:]), start=1):
        ffmpeg_cmd = [
            'ffmpeg',
            '-ss', start]
        if end is not None:
            ffmpeg_cmd.extend(['-to', end])

        out_path = output_dir / name_template.format(i)

        ffmpeg_cmd.extend([
            '-i', input_file,
            '-c:a', 'flac',
            '-movflags', '+faststart',
            str(out_path)])

        subprocess.check_call(ffmpeg_cmd)


if __name__ == '__main__':
    cli(complete_var='_RIP_PY_COMPLETE')