ncmr-songs-videos/scripts/srt2ass.py

"""Convert an SRT subtitle file to ASS (Advanced SubStation Alpha).

Stdlib only. Handles the constructs that actually appear in song subtitles:
multi-line cues, ``{\\an8}``-style override blocks passed through untouched,
and the SRT quirks (CRLF, BOM, cue numbers, blank-line separators).

Usage::

    python3 scripts/srt2ass.py song.srt song.ass
    python3 scripts/srt2ass.py lyrics/*.srt --out-dir build/  # batch
    python3 scripts/srt2ass.py song.srt --style-font "Noto Sans" --style-size 64
"""

from __future__ import annotations

import argparse
import os
import re
import sys
from dataclasses import dataclass

if __package__ in (None, ""):
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from _common import (  # type: ignore[import-not-found]
        Die,
        Report,
        add_common_arguments,
        glob_files,
        main_wrapper,
    )
else:  # pragma: no cover
    from ._common import (
        Die,
        Report,
        add_common_arguments,
        glob_files,
        main_wrapper,
    )


# --------------------------------------------------------------------------- #
# Timestamps
# --------------------------------------------------------------------------- #

_SRT_TIME_RE = re.compile(
    r"(?P<h>\d{1,2}):(?P<m>\d{2}):(?P<s>\d{2})[,.](?P<ms>\d{1,3})"
)


def parse_srt_time(stamp: str) -> float:
    """``00:01:02,500`` (or ``00:01:02.5``) to seconds."""
    match = _SRT_TIME_RE.search(stamp.strip())
    if not match:
        raise Die(f"unparseable SRT timestamp: {stamp!r}")

    ms = match.group("ms").ljust(3, "0")[:3]

    return (
        int(match.group("h")) * 3600
        + int(match.group("m")) * 60
        + int(match.group("s"))
        + int(ms) / 1000
    )


def format_ass_time(seconds: float) -> str:
    """Seconds to ASS ``H:MM:SS.cc`` (centiseconds, single-digit hour)."""
    if seconds < 0:
        seconds = 0.0

    total_cs = int(round(seconds * 100))
    hours, remainder = divmod(total_cs, 360_000)
    minutes, remainder = divmod(remainder, 6_000)
    secs, centis = divmod(remainder, 100)

    return f"{hours}:{minutes:02d}:{secs:02d}.{centis:02d}"


# --------------------------------------------------------------------------- #
# Parsing
# --------------------------------------------------------------------------- #


@dataclass
class Cue:
    index: int
    start: float
    end: float
    lines: list[str]

    def text(self, line_sep: str = r"\N") -> str:
        return line_sep.join(self.lines)


def strip_bom(text: str) -> str:
    return text[1:] if text.startswith("\ufeff") else text


def parse_srt(content: str) -> list[Cue]:
    """Parse SRT text into cues. Tolerates CRLF, BOM, and missing cue numbers."""
    content = strip_bom(content).replace("\r\n", "\n").replace("\r", "\n")
    blocks = re.split(r"\n\s*\n", content)

    cues: list[Cue] = []
    auto_index = 0

    for block in blocks:
        lines = block.strip("\n").split("\n")
        if not lines or not any(line.strip() for line in lines):
            continue

        index: int | None = None
        if lines[0].strip().isdigit():
            index = int(lines[0].strip())
            lines = lines[1:]

        if not lines:
            continue

        if "-->" not in lines[0]:
            # Not a cue (stray header, WEBVTT marker, ...). Skip it quietly.
            continue

        start_raw, _, end_raw = lines[0].partition("-->")
        start = parse_srt_time(start_raw)
        end = parse_srt_time(end_raw)

        body = [line for line in lines[1:] if line.strip()]
        if not body:
            # A cue with a timestamp but no text: pure gap in the lyrics. Emitting
            # it would produce an empty subtitle block, so drop it — the previous
            # cue simply holds until the next one.
            continue

        auto_index += 1
        cues.append(Cue(index=index or auto_index, start=start, end=end, lines=body))

    return cues


# --------------------------------------------------------------------------- #
# ASS output
# --------------------------------------------------------------------------- #


#: ASS uses US-ASCII-with-code-page-1252 in practice; karaoke/Japanese needs 1.
ASS_CODEPAGES = {
    "ascii": 0,
    "latin1": 0,
    "cp1252": 0,
    "shift_jis": 1,
    "sjis": 1,
    "gbk": 134,
    "utf-8": 1,
    "utf8": 1,
}

ASS_HEADER = """[Script Info]
; Generated by scripts/srt2ass.py — edit the source SRT, not this file.
Title: {title}
ScriptType: v4.00+
WrapStyle: {wrap_style}
ScaledBorderAndShadow: yes
YCbCr Matrix: TV.709
PlayResX: {width}
PlayResY: {height}

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,{font},{size},{primary},{secondary},{outline_colour},{back},{bold},{italic},0,0,100,100,{spacing},0,1,{border},{shadow},{alignment},{margin_l},{margin_r},{margin_v},{encoding}

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
"""


@dataclass
class AssStyle:
    """The subset of the Default style we expose on the CLI."""

    font: str = "Arial"
    size: int = 48
    primary: str = "&H00FFFFFF"  # white, fully opaque
    secondary: str = "&H000000FF"
    outline_colour: str = "&H00000000"
    back: str = "&H00000000"
    bold: int = -1
    italic: int = 0
    spacing: float = 0.0
    border: float = 2.0
    shadow: float = 1.0
    alignment: int = 2  # 2 = bottom centre
    margin_l: int = 40
    margin_r: int = 40
    margin_v: int = 60
    encoding: int = 1


def render_ass(
    cues: list[Cue],
    *,
    style: AssStyle,
    width: int,
    height: int,
    title: str,
    wrap_style: int = 0,
) -> str:
    header = ASS_HEADER.format(
        title=title.replace("\n", " ") or "Untitled",
        wrap_style=wrap_style,
        width=width,
        height=height,
        font=style.font,
        size=style.size,
        primary=style.primary,
        secondary=style.secondary,
        outline_colour=style.outline_colour,
        back=style.back,
        bold=style.bold,
        italic=style.italic,
        spacing=style.spacing,
        border=style.border,
        shadow=style.shadow,
        alignment=style.alignment,
        margin_l=style.margin_l,
        margin_r=style.margin_r,
        margin_v=style.margin_v,
        encoding=style.encoding,
    )

    events = []
    for cue in cues:
        text = cue.text()
        # ASS treats commas as field separators in the Text column, but only the
        # first 9 commas matter — ffmpeg/libass handles commas inside Text fine as
        # long as the Format line lists Text last, which it does above.
        events.append(
            "Dialogue: 0,{start},{end},Default,,0,0,0,,{text}".format(
                start=format_ass_time(cue.start),
                end=format_ass_time(cue.end),
                text=text,
            )
        )

    return header + "\n".join(events) + "\n"


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="srt2ass",
        description="Convert SRT subtitles to ASS, ready for ffmpeg's ass= filter.",
        epilog=(
            "Then bake with:  ffmpeg -vf \"ass=out.ass\" -i in.mp4 ...  "
            "(or scripts/ffmpeg_tools.py bake)"
        ),
    )
    add_common_arguments(parser)

    parser.add_argument("inputs", nargs="+", metavar="INPUT", help="SRT file(s) to convert")
    parser.add_argument(
        "-o",
        "--output",
        metavar="FILE",
        help="output .ass path (single input only; default: alongside the input)",
    )
    parser.add_argument(
        "--out-dir",
        metavar="DIR",
        help="write every output into this directory (created if missing)",
    )

    style = parser.add_argument_group("style")
    style.add_argument("--style-font", default="Arial", help="font name (default: Arial)")
    style.add_argument("--style-size", type=int, default=48, help="font size (default: 48)")
    style.add_argument(
        "--style-alignment",
        type=int,
        default=2,
        choices=range(1, 10),
        metavar="N",
        help="ASS numpad alignment 1-9 (default: 2, bottom centre)",
    )
    style.add_argument("--style-margin-v", type=int, default=60, help="vertical margin (default: 60)")
    style.add_argument("--style-outline", type=float, default=2.0, help="outline width (default: 2)")
    style.add_argument("--style-shadow", type=float, default=1.0, help="shadow depth (default: 1)")
    style.add_argument("--style-bold", action=argparse.BooleanOptionalAction, default=True)
    style.add_argument("--style-italic", action=argparse.BooleanOptionalAction, default=False)

    canvas = parser.add_argument_group("canvas")
    canvas.add_argument("--width", type=int, default=1920, help="PlayResX (default: 1920)")
    canvas.add_argument("--height", type=int, default=1080, help="PlayResY (default: 1080)")
    canvas.add_argument(
        "--title",
        default=None,
        help="ASS Title field (default: the output filename)",
    )

    return parser


def convert_one(src: str, dest: str, args: argparse.Namespace, report: Report) -> tuple[int, int]:
    try:
        with open(src, "r", encoding="utf-8-sig", errors="replace") as handle:
            content = handle.read()
    except OSError as exc:
        raise Die(f"cannot read {src}: {exc.strerror}") from exc

    cues = parse_srt(content)
    if not cues:
        report.warn(f"{src}: no cues parsed — is this really an SRT file?")

    style = AssStyle(
        font=args.style_font,
        size=args.style_size,
        alignment=args.style_alignment,
        margin_v=args.style_margin_v,
        border=args.style_outline,
        shadow=args.style_shadow,
        bold=-1 if args.style_bold else 0,
        italic=-1 if args.style_italic else 0,
    )

    title = args.title or os.path.splitext(os.path.basename(dest))[0]
    document = render_ass(
        cues,
        style=style,
        width=args.width,
        height=args.height,
        title=title,
    )

    if args.dry_run:
        report.info(f"[dry-run] would write {dest} ({len(cues)} cues)")
        return len(cues), 0

    os.makedirs(os.path.dirname(os.path.abspath(dest)), exist_ok=True)
    with open(dest, "w", encoding="utf-8", newline="\n") as handle:
        handle.write(document)

    report.ok(f"{src} -> {dest} ({len(cues)} cues)")

    return len(cues), 0


@main_wrapper
def main() -> int:
    args = build_parser().parse_args()
    report = Report(quiet=args.quiet)

    if args.output and len(args.inputs) > 1:
        raise Die("--output takes a single input; use --out-dir for batches")

    total = 0
    for src in args.inputs:
        if not os.path.isfile(src):
            raise Die(f"no such file: {src}")

        stem = os.path.splitext(os.path.basename(src))[0]

        if args.output:
            dest = args.output
        elif args.out_dir:
            dest = os.path.join(args.out_dir, stem + ".ass")
        else:
            dest = os.path.splitext(src)[0] + ".ass"

        cues, _ = convert_one(src, dest, args, report)
        total += cues

    if len(args.inputs) > 1:
        report.info(f"{total} cues across {len(args.inputs)} file(s)")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())