youtube_abschrift_korrigieren/extract_text.py

#!/usr/bin/env python3
"""Lädt die Untertitel eines YouTube-Videos und speichert ihren Text als Datei
(yt_untertitel.py + vtt2txt.py in einem Schritt)."""
from pathlib import Path
import argparse
import shutil
import sys
import tempfile
import textwrap

from vtt2txt import (
    add_time_options, join_segments, range_argv, resolve_time_ranges, split_segments,
    time_suffix, vtt_to_text,
)
from yt_untertitel import (
    YT_DLP, choose_track, default_name, download, fetch_info, no_track_message,
    original_language, work_folder, write_source,
)

FLOW_WIDTH = 80  # Zeilenbreite für --fliesstext


def parse_args():
    parser = argparse.ArgumentParser(
        description="Lädt die Untertitel eines YouTube-Videos und speichert den Text."
    )
    parser.add_argument("-u", "--url", required=True, help="YouTube-URL des Videos")
    parser.add_argument(
        "-l", "--lang",
        help="Sprache der Untertitel, z. B. de (Standard: Sprache des Original-Audios)",
    )
    parser.add_argument(
        "-o", "--output", type=Path,
        help="Ausgabedatei (Standard: abschriften/<Name>/<Name>.txt mit "
        "<Name> = <Kanal>_<Titel>_<Datum>, mit -z zusätzlich die Bereiche)",
    )
    parser.add_argument(
        "-f", "--fliesstext", action="store_true",
        help=f"Fließtext, auf {FLOW_WIDTH} Zeichen umbrochen, statt Untertitelzeilen",
    )
    add_time_options(parser)
    args = parser.parse_args(range_argv(sys.argv[1:]))
    resolve_time_ranges(parser, args)

    if args.output is not None and args.output.is_dir():
        parser.error(f"-o/--output ist ein Verzeichnis: {args.output}")
    if shutil.which(YT_DLP) is None:
        parser.error(f"{YT_DLP} nicht gefunden (Installation: pip install yt-dlp).")
    return args


def fetch_text(url, lang=None, ranges=(), exact=False):
    """Download the subtitles of a video; return (video info, text).

    Without `lang` the language of the original audio is used. With time
    `ranges` (seconds) the text has one block per range under a header with
    its actual times, see vtt2txt.vtt_to_text.
    """
    info = fetch_info(url)
    print(f"Video: {info.get('title', '?')}")

    lang = lang or original_language(info)
    chosen = choose_track(info, lang) if lang else None
    if chosen is None:
        raise SystemExit(no_track_message(info, lang))
    track, is_auto = chosen
    kind = "automatische" if is_auto else "manuelle"
    print(f"Lade {kind} Untertitel „{track}“ …", flush=True)

    # The VTT file is only an intermediate step
    with tempfile.TemporaryDirectory() as tmp:
        target_base = Path(tmp) / "untertitel"
        download(url, track, is_auto, target_base)
        vtt = Path(f"{target_base}.{track}.vtt")
        if not vtt.is_file():
            raise SystemExit("yt-dlp hat keine Untertitel-Datei geschrieben.")
        text = vtt_to_text(vtt.read_text(encoding="utf-8-sig"), ranges, exact)

    if not text:
        if ranges:
            raise SystemExit("Die Untertitel enthalten im gewählten Zeitraum keinen Text.")
        raise SystemExit("Die Untertitel enthalten keinen Text.")
    return info, text


def main():
    args = parse_args()
    info, text = fetch_text(args.url, args.lang, args.bereiche, args.exakt)
    suffix = time_suffix(args.bereiche)
    name = default_name(info)
    output = args.output or work_folder(name) / f"{name}{suffix}.txt"
    segments = split_segments(text)

    if args.fliesstext:
        # Each block on its own, headers stay lines of their own; never split
        # words, not even at hyphens ("EU-Gipfel")
        segments = [
            (header, textwrap.fill(
                " ".join(body.split()), width=FLOW_WIDTH,
                break_long_words=False, break_on_hyphens=False,
            ))
            for header, body in segments
        ]
        text = join_segments(segments)

    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(text + "\n", encoding="utf-8")
    write_source(info, output.with_suffix(".quelle.json"))
    print(
        f"Gespeichert: {output} "
        f"({len(text.splitlines())} Zeilen, "
        f"{sum(len(body.split()) for _, body in segments)} Wörter)"
    )


if __name__ == "__main__":
    main()