ncmr-songs-videos/scripts/_common.py

"""NCMR songs-videos helper toolchain — Python, stdlib only.

Shared plumbing for the repo's helper scripts. Every script in this package is
runnable as a plain file (``python3 scripts/bom_gen.py``) and as a module
(``python3 -m scripts.bom_gen``); both work because this package is imported
either way without needing its parent directory on ``sys.path``.

Nothing outside the standard library is imported anywhere in this package.
"""

import argparse
import os
import re
import shutil
import subprocess
import sys
from dataclasses import dataclass, field
from typing import Callable, Iterable, Sequence

__all__ = [
    "Die",
    "VERSION",
    "Report",
    "find_ffmpeg",
    "glob_files",
    "human_size",
    "report",
    "run",
    "subparsers_metavar",
    "text_table",
]

VERSION = "0.2.0"


class Die(Exception):
    """Fatal, user-facing error. Caught in ``main()`` and printed to stderr."""


# --------------------------------------------------------------------------- #
# Console output
# --------------------------------------------------------------------------- #


def _colour_enabled(stream=None) -> bool:
    stream = stream or sys.stdout

    if os.environ.get("NO_COLOR"):
        return False
    if os.environ.get("FORCE_COLOR"):
        return True

    return hasattr(stream, "isatty") and stream.isatty()


def _style(text: str, code: str) -> str:
    if not _colour_enabled():
        return text

    return f"\033[{code}m{text}\033[0m"


@dataclass
class Report:
    """Prints user-facing output, with optional ANSI colour when on a TTY."""

    quiet: bool = False

    def _emit(self, message: str, *, err: bool = False) -> None:
        stream = sys.stderr if err else sys.stdout
        print(message, file=stream)

    def info(self, message: str = "") -> None:
        """Ordinary progress output."""
        self._emit(message)

    def ok(self, message: str) -> None:
        self._emit(_style(f"✓ {message}", "32"))

    def warn(self, message: str) -> None:
        self._emit(_style(f"warning: {message}", "33"), err=True)

    def error(self, message: str) -> None:
        self._emit(_style(f"error: {message}", "31"), err=True)

    def verbose(self, message: str) -> None:
        if not self.quiet:
            self._emit(_style(f"  {message}", "2"))


#: Default reporter used when a caller doesn't supply one.
report = Report()


# --------------------------------------------------------------------------- #
# Formatting helpers
# --------------------------------------------------------------------------- #

_SIZE_UNITS = ("B", "KiB", "MiB", "GiB", "TiB")

# Characters that make a key unquoted-safe in YAML and legible in a table.
_SAFE_KEY_RE = re.compile(r"^[A-Za-z0-9_.\-/]+$")


def human_size(num_bytes: int) -> str:
    """Format a byte count as ``1.5 MiB`` / ``12 B``."""
    value = float(num_bytes)

    for unit in _SIZE_UNITS[:-1]:
        if value < 1024:
            break
        value /= 1024
    else:
        unit = _SIZE_UNITS[-1]

    if unit == "B":
        return f"{int(value)} B"

    return f"{value:.1f} {unit}"


def text_table(headers: Sequence[str], rows: Iterable[Sequence[str]]) -> str:
    """Render a simple whitespace-aligned table (no boxes, pipe-friendly)."""
    body = [[str(cell) for cell in row] for row in rows]
    widths = [len(str(h)) for h in headers]

    for row in body:
        for i, cell in enumerate(row):
            if i < len(widths):
                widths[i] = max(widths[i], len(cell))

    def line(cells: Sequence[str]) -> str:
        return "  ".join(cell.ljust(widths[i]) for i, cell in enumerate(cells)).rstrip()

    out = [line([str(h) for h in headers])]
    out.append("  ".join("-" * w for w in widths).rstrip())

    for row in body:
        out.append(line(row))

    return "\n".join(out)


def glob_files(directory: str, suffix: str) -> list[str]:
    """Sorted list of files directly inside *directory* whose name ends in *suffix*.

    Returns an empty list when the directory does not exist, so callers can treat
    "no such folder" and "folder is empty" identically.
    """
    try:
        entries = os.listdir(directory)
    except (FileNotFoundError, NotADirectoryError):
        return []

    suffix = suffix.lower()
    found = [
        os.path.join(directory, name)
        for name in entries
        if name.lower().endswith(suffix) and os.path.isfile(os.path.join(directory, name))
    ]

    return sorted(found)


# --------------------------------------------------------------------------- #
# Subprocess plumbing
# --------------------------------------------------------------------------- #


def run(
    cmd: Sequence[str],
    *,
    report: Report | None = None,
    check: bool = True,
    cwd: str | None = None,
    dry_run: bool = False,
) -> subprocess.CompletedProcess:
    """Run *cmd* with a shell-free argument list, echoing it first.

    ``report`` is the reporter used for the echo; the helper deliberately takes a
    reporter argument rather than closing over the module-level default so that a
    ``--quiet`` CLI cannot accidentally echo commands.
    """
    log = report or globals()["report"]
    printable = " ".join(shlex_quote(part) for part in cmd)

    if dry_run:
        log.info(_style(f"[dry-run] {printable}", "36"))
        return subprocess.CompletedProcess(list(cmd), 0)

    log.info(_style(f"$ {printable}", "2"))

    # ffmpeg prints its whole build configuration as a banner unless told not to.
    # Keep it quiet by default, and keep stderr so a failure can be explained.
    full_cmd = list(cmd)
    if os.path.basename(full_cmd[0]) in ("ffmpeg", "ffprobe"):
        full_cmd[1:1] = ["-hide_banner"]

    try:
        result = subprocess.run(full_cmd, cwd=cwd, capture_output=True, text=True)
    except FileNotFoundError as exc:
        raise Die(f"command not found: {cmd[0]}") from exc

    if result.stdout:
        sys.stdout.write(result.stdout)
    if result.stderr:
        # ffmpeg writes progress to stderr; pass it through so the user still sees
        # the encode happen, but drop the wall of banner text.
        sys.stderr.write(result.stderr)

    if check and result.returncode != 0:
        detail = _last_meaningful_line(result.stderr)
        suffix = f": {detail}" if detail else ""
        raise Die(f"{os.path.basename(cmd[0])} failed (exit {result.returncode}){suffix}")

    return result


def _last_meaningful_line(text: str | None) -> str:
    """Pick the most useful line out of ffmpeg's stderr to show in an error."""
    if not text:
        return ""

    for line in reversed(text.strip().splitlines()):
        line = line.strip()
        if not line:
            continue
        # Progress lines start with 'frame=' and overwrite each other.
        if line.startswith(("frame=", "size=")) or "speed=" in line:
            continue

        return line

    return ""


def shlex_quote(part: str) -> str:
    """Quote one argv element for display. Local copy keeps the module stdlib-only."""
    if not part:
        return "''"
    if re.fullmatch(r"[A-Za-z0-9_@%+=:,./\-]+", part):
        return part

    return "'" + part.replace("'", "'\"'\"'") + "'"


def find_ffmpeg(override: str | None = None) -> str:
    """Locate the ``ffmpeg`` binary, or raise :class:`Die` with a useful message."""
    if override:
        if os.path.sep in override:
            if not os.access(override, os.X_OK):
                raise Die(f"--ffmpeg {override} is not an executable file")
            return override

        return override  # a bare name; let the OS resolve it on $PATH

    found = shutil.which("ffmpeg")
    if not found:
        raise Die("ffmpeg not found on $PATH — install it or pass --ffmpeg /path/to/ffmpeg")

    return found


# --------------------------------------------------------------------------- #
# argparse ergonomics
# --------------------------------------------------------------------------- #


def subparsers_metavar(parser: argparse.ArgumentParser, name: str = "mode") -> None:
    """Stop argparse from printing an unreadable ``{a,b,c}`` in help output.

    Call *after* every subparser has been registered.
    """
    for action in parser._actions:
        if isinstance(action, argparse._SubParsersAction):
            action.metavar = name
            action.choices = None  # type: ignore[assignment]  # keep choices working


def add_common_arguments(parser: argparse.ArgumentParser, *, suppress_defaults: bool = False) -> None:
    """Switches every tool in the toolbox understands.

    ``suppress_defaults`` is used when the same switches are also registered on
    subparsers: it stops the subparser copy from stamping its default over a value
    the user gave before the mode word.
    """
    default_suppress = argparse.SUPPRESS if suppress_defaults else False

    parser.add_argument(
        "-q",
        "--quiet",
        action="store_true",
        default=default_suppress,
        help="suppress progress output (errors still go to stderr)",
    )
    parser.add_argument(
        "-n",
        "--dry-run",
        action="store_true",
        default=default_suppress,
        help="print the work that would happen without doing it",
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"%(prog)s {VERSION}",
    )


def add_ffmpeg_argument(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--ffmpeg",
        metavar="PATH",
        default=None,
        help="ffmpeg binary to use (default: whichever is on $PATH)",
    )


@dataclass
class ReportMixin:
    """Mixin for subcommand classes that need a reporter and a dry-run flag."""

    report: Report = field(default_factory=Report)
    dry_run: bool = False


def main_wrapper(fn: Callable[[], int | None]) -> Callable[[], int]:
    """Decorator turning ``Die`` exceptions into a clean exit status."""

    def wrapper() -> int:
        try:
            return fn() or 0
        except Die as exc:
            report.error(str(exc))
            return 1
        except KeyboardInterrupt:
            report.error("interrupted")
            return 130

    return wrapper