GxEPD2-esphome/tools/capture.py

#!/usr/bin/env python3
"""One-command panel-descriptor capture.

Resolves a GxEPD2 driver class to its vendored source, builds and runs a host
harness that *observes* the driver, and prints a candidate EpdPanelDescriptor
(the body of a `components/gxepd2/panel_<model>.h`). No ESP toolchain needed -
just a host C++ compiler. See tools/gxepd2_harness/README.md and claude/DESIGN.md.

    tools/capture.py GxEPD2_154_Z90c             # candidate -> stdout
    tools/capture.py GxEPD2_154_Z90c -o panel.h  # candidate -> file
    tools/capture.py --list                      # capturable driver classes
    tools/capture.py GxEPD2_154_Z90c --debug     # also dump the capture structure

The output is a *candidate*: review it, fill any TODO fields the diag passes
own (bits_per_pixel/palette for packed panels; plane_width/row_interleave for
dual-region), save it as components/gxepd2/panel_<model>.h, add it to
display.py's MODELS, compile, and bench-verify before marking it verified.
"""

import argparse
import difflib
import json
import pathlib
import re
import subprocess
import sys
import tempfile

ROOT = pathlib.Path(__file__).resolve().parent.parent  # code/
CATALOG = ROOT / "tools" / "gxepd2_catalog" / "catalog.json"
VSRC = ROOT / "third_party" / "GxEPD2" / "src"
HARNESS = ROOT / "tools" / "gxepd2_harness"
DISPLAY_PY = ROOT / "components" / "gxepd2" / "display.py"
PREFIX = "GxEPD2_"


def macro_of(model: str) -> str:
    return "PANEL_" + model[len(PREFIX):].upper()


def filename_of(model: str) -> str:
    return "panel_" + model[len(PREFIX):].lower() + ".h"


def shipped_models() -> set:
    """Model keys already wired into display.py's MODELS, to flag in --list."""
    try:
        text = DISPLAY_PY.read_text()
    except OSError:
        return set()
    return set(re.findall(r'"(GxEPD2_[^"]+)":\s*\(', text))


def load_catalog() -> dict:
    if not CATALOG.exists():
        sys.exit(f"catalog not found: {CATALOG}\n(run tools/gxepd2_catalog/generate_catalog.py)")
    return json.loads(CATALOG.read_text())


def capturable(entry: dict) -> bool:
    # The emitter instantiates Driver(cs, dc, rst, busy) - only the standard
    # 4-arg constructor works without special-casing.
    return entry.get("ctor_args") == 4 and (VSRC / entry["source"]).exists()


def do_list(catalog: dict) -> None:
    shipped = shipped_models()
    rows = []
    for model, entry in sorted(catalog.items()):
        if not capturable(entry):
            continue
        mark = "shipped" if model in shipped else "-"
        rows.append((model, mark, entry.get("comment", "")))
    w = max(len(m) for m, _, _ in rows)
    print(f"{len(rows)} capturable driver classes (ctor_args=4, source present):\n")
    for model, mark, comment in rows:
        print(f"  {model:<{w}}  {mark:<7}  {comment}")


def resolve_source(model: str, catalog: dict):
    """(source relpath, ctor_args|None). Catalog first (it has ctor_args), else
    glob the vendored tree - so a driver the fork added but GxEPD2's example
    headers don't list (e.g. the F52) is still capturable."""
    if model in catalog:
        e = catalog[model]
        return e["source"], e.get("ctor_args")
    hits = list(VSRC.glob(f"**/{model}.cpp"))
    if len(hits) == 1:
        return str(hits[0].relative_to(VSRC)), None
    if len(hits) > 1:
        sys.exit(f"{model}: multiple sources found: {[str(h) for h in hits]}")
    return None, None


def capture(model: str, catalog: dict, debug: bool, cxx: str, keep: bool) -> str:
    src_rel, ctor_args = resolve_source(model, catalog)
    if src_rel is None:
        names = set(catalog) | {p.stem for p in VSRC.glob("**/GxEPD2_*.cpp")}
        hint = difflib.get_close_matches(model, names, n=3)
        msg = f"unknown model: {model} (no {model}.cpp under {VSRC})"
        if hint:
            msg += "\ndid you mean: " + ", ".join(hint)
        sys.exit(msg)
    if ctor_args not in (None, 4):
        sys.exit(f"{model}: non-standard constructor (ctor_args={ctor_args}); "
                 "the emitter only handles the 4-arg GxEPD2_EPD(cs,dc,rst,busy) form.")
    src = VSRC / src_rel
    hdr_rel = src_rel[:-4] + ".h"  # epd3c/GxEPD2_154_Z90c.h
    if not (VSRC / hdr_rel).exists():
        sys.exit(f"driver header not found: {VSRC / hdr_rel}")

    macro = macro_of(model)
    with tempfile.TemporaryDirectory() as td:
        tdp = pathlib.Path(td)
        main_cpp = tdp / "emit_main.cpp"
        main_cpp.write_text(
            f'#include "emit_descriptor.h"\n'
            f'#include "{hdr_rel}"\n'
            f'int main() {{ emit::emit_descriptor<{model}>("{model}", "{macro}"); return 0; }}\n'
        )
        binary = tdp / "emit_main"
        compile_cmd = [
            cxx, "-std=c++17",
            "-I", str(HARNESS / "spy"), "-I", str(VSRC), "-I", str(HARNESS),
            str(main_cpp), str(VSRC / "GxEPD2_EPD.cpp"), str(src),
            "-o", str(binary),
        ]
        print(f"[capture] {model}  <-  third_party/GxEPD2/src/{src_rel}", file=sys.stderr)
        cc = subprocess.run(compile_cmd, capture_output=True, text=True)
        if cc.returncode != 0:
            sys.exit(f"compile failed:\n{cc.stderr}")
        env = {"EMIT_DEBUG": "1"} if debug else None
        run = subprocess.run([str(binary)], capture_output=True, text=True, env=env)
        if keep:
            kept = HARNESS / f"run_{model[len(PREFIX):].lower()}_emit.cpp"
            kept.write_text(main_cpp.read_text())
            print(f"[capture] kept emit main: {kept}", file=sys.stderr)
    if debug and run.stderr:
        print(run.stderr, file=sys.stderr)
    if run.returncode != 0:
        sys.exit(f"capture run failed (exit {run.returncode}):\n{run.stdout}\n{run.stderr}")

    marker = "// GENERATED by emit_descriptor"
    idx = run.stdout.find(marker)
    if idx < 0:
        sys.exit(f"no descriptor in output:\n{run.stdout}")
    descriptor = run.stdout[idx:]
    print(f"[capture] ok -> save as components/gxepd2/{filename_of(model)}, "
          f"add \"{model}\" to display.py MODELS", file=sys.stderr)
    return descriptor


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("model", nargs="?", help="GxEPD2 driver class, e.g. GxEPD2_154_Z90c")
    ap.add_argument("-o", "--out", help="write the candidate descriptor to this file (default: stdout)")
    ap.add_argument("--list", action="store_true", help="list capturable driver classes and exit")
    ap.add_argument("--debug", action="store_true", help="dump the image-stage command/burst structure to stderr")
    ap.add_argument("--keep", action="store_true", help="keep the generated emit main under tools/gxepd2_harness/")
    ap.add_argument("--cxx", default="g++", help="C++ compiler (default: g++)")
    args = ap.parse_args()

    catalog = load_catalog()
    if args.list:
        do_list(catalog)
        return
    if not args.model:
        ap.error("a model is required (or use --list)")
    descriptor = capture(args.model, catalog, args.debug, args.cxx, args.keep)
    if args.out:
        pathlib.Path(args.out).write_text(descriptor)
        print(f"[capture] wrote {args.out}", file=sys.stderr)
    else:
        sys.stdout.write(descriptor)


if __name__ == "__main__":
    main()