nordisk/facett

py/scripts/regen_protos.py

#!/usr/bin/env python3
"""Regenerate this package's gRPC stubs FROM THE SOURCE OF TRUTH.

    python scripts/regen_protos.py [--check]

Two wires, two different rules:

* **korp** — `korp/proto/korp.proto` in the sibling checkout is the ONE copy. This
  package used to carry a hand-maintained duplicate of it; that duplicate drifted
  three RPCs behind the server before anyone noticed. It is gone. Generate from
  the sibling (override the location with `KORP_PROTO`).
* **saga** — `facett_demo/proto/saga.proto` is a deliberate wire-compatible
  SUBSET of `edda/crates/nornir-wire/proto/nornir.proto` (one service out of
  twenty), so it stays checked in. `tests/test_saga.py` is what keeps it honest.

`--check` regenerates into a temp dir and diffs, so CI can assert the checked-in
stubs are what this script would produce. Exit 1 on any difference.
"""

from __future__ import annotations

import argparse
import filecmp
import os
import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

HERE = Path(__file__).resolve().parent
PKG_PROTO = HERE.parent / "facett_demo" / "proto"

# Walked, not indexed: `HERE.parents[2]` raises `IndexError` when the package is
# mounted shallow (`/app` in the demo's container), and this script has to run THERE
# — no python runs on the host. `KORP_PROTO` still wins.
_KORP_REL = Path("korp") / "proto" / "korp.proto"
_KORP_CANDIDATES = [parent / _KORP_REL for parent in HERE.parents]


def korp_proto() -> Path:
    env = os.environ.get("KORP_PROTO")
    if env:
        p = Path(env)
        if not p.is_file():
            sys.exit(f"KORP_PROTO={env!r} is not a file")
        return p
    for c in _KORP_CANDIDATES:
        if c.is_file():
            return c
    sys.exit(
        "korp's proto was not found. Looked at:\n  "
        + "\n  ".join(str(c) for c in _KORP_CANDIDATES)
        + "\nCheck out the sibling `korp` repo, or set KORP_PROTO."
    )


def _protoc(proto: Path, out: Path) -> None:
    out.mkdir(parents=True, exist_ok=True)
    cmd = [
        sys.executable, "-m", "grpc_tools.protoc",
        f"-I{proto.parent}",
        f"--python_out={out}",
        f"--grpc_python_out={out}",
        proto.name,
    ]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode:
        sys.exit(f"protoc failed for {proto}:\n{r.stderr}")
    # grpc_tools emits a top-level `import foo_pb2`; inside a package that must be
    # a relative import or the stub is unimportable.
    stub = out / f"{proto.stem}_pb2_grpc.py"
    text = stub.read_text()
    fixed = re.sub(rf"^import {proto.stem}_pb2 as", f"from . import {proto.stem}_pb2 as",
                   text, flags=re.M)
    if fixed != text:
        stub.write_text(fixed)


def generate(out: Path) -> list[str]:
    stems = []
    _protoc(korp_proto(), out)
    stems.append("korp")
    saga = PKG_PROTO / "saga.proto"
    if saga.is_file():
        _protoc(saga, out)
        stems.append("saga")
    return stems


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--check", action="store_true",
                    help="verify the checked-in stubs match, do not write")
    args = ap.parse_args()

    if not args.check:
        stray = PKG_PROTO / "korp.proto"
        if stray.exists():
            stray.unlink()
            print(f"removed the hand copy {stray}")
        stems = generate(PKG_PROTO)
        print(f"regenerated {', '.join(s + '_pb2[_grpc].py' for s in stems)} in {PKG_PROTO}")
        print(f"korp source of truth: {korp_proto()}")
        return 0

    with tempfile.TemporaryDirectory() as td:
        tmp = Path(td)
        stems = generate(tmp)
        bad = []
        for stem in stems:
            for suffix in ("_pb2.py", "_pb2_grpc.py"):
                name = stem + suffix
                a, b = PKG_PROTO / name, tmp / name
                if not a.exists() or not filecmp.cmp(a, b, shallow=False):
                    bad.append(name)
        if (PKG_PROTO / "korp.proto").exists():
            bad.append("korp.proto (a hand copy of another repo's wire)")
        if bad:
            print("STALE — these do not match what the source protos generate:", file=sys.stderr)
            for n in bad:
                print(f"  {n}", file=sys.stderr)
            print("Run: python scripts/regen_protos.py", file=sys.stderr)
            return 1
        print("stubs are in step with the source protos")
        return 0


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