cmvr_edge_ai/scripts/generate_cmvr_stubs.py
2026-07-20 16:59:37 +08:00

118 lines
3.8 KiB
Python

#!/usr/bin/env python3
"""Generate Python protobuf/gRPC bindings from the cmvr-es source tree."""
from __future__ import annotations
import argparse
import platform
import subprocess
import sys
from pathlib import Path
DEFAULT_CMVR_ES_ROOT = Path("/home/xtkuang/Projects/cmvr/cmvr-es")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--cmvr-es-root",
type=Path,
default=DEFAULT_CMVR_ES_ROOT,
help="cmvr-es checkout containing protos/ and output/bin/",
)
parser.add_argument(
"--output",
type=Path,
default=Path("src"),
help="output root; generated imports require a top-level cmvr package",
)
parser.add_argument(
"--portable",
action="store_true",
help="use python -m grpc_tools.protoc instead of cmvr-es bundled tools",
)
return parser.parse_args()
def build_command(args: argparse.Namespace, proto_files: list[Path]) -> list[str]:
proto_root = args.cmvr_es_root.resolve() / "protos"
output = args.output.resolve()
relative_files = [str(path.relative_to(proto_root)) for path in proto_files]
if args.portable:
try:
import grpc_tools
except ImportError as exc:
raise RuntimeError(
"portable generation requires grpcio-tools; install the dev extra"
) from exc
well_known = Path(grpc_tools.__file__).resolve().parent / "_proto"
return [
sys.executable,
"-m",
"grpc_tools.protoc",
f"-I{proto_root}",
f"-I{well_known}",
f"--python_out={output}",
f"--pyi_out={output}",
f"--grpc_python_out={output}",
*relative_files,
]
protoc = args.cmvr_es_root.resolve() / "output/bin/protoc"
grpc_plugin = args.cmvr_es_root.resolve() / "output/bin/grpc_python_plugin"
if not protoc.is_file() or not grpc_plugin.is_file():
raise FileNotFoundError(
"cmvr-es protobuf tools are missing; build cmvr-es first or pass --portable"
)
well_known = _bundled_well_known_proto_root(args.cmvr_es_root.resolve())
return [
str(protoc),
f"-I{proto_root}",
f"-I{well_known}",
f"--python_out={output}",
f"--pyi_out={output}",
f"--grpc_python_out={output}",
f"--plugin=protoc-gen-grpc_python={grpc_plugin}",
*relative_files,
]
def _bundled_well_known_proto_root(cmvr_es_root: Path) -> Path:
architecture = (
"arm" if platform.machine().lower() in {"aarch64", "arm64"} else "x86"
)
candidates = (
cmvr_es_root / "dependency" / architecture / "third_party/grpc/v1.76.0/include",
cmvr_es_root / "dependency/x86/third_party/grpc/v1.76.0/include",
cmvr_es_root / "dependency/arm/third_party/grpc/v1.76.0/include",
)
for candidate in candidates:
if (candidate / "google/protobuf/timestamp.proto").is_file():
return candidate
raise FileNotFoundError(
"google/protobuf/timestamp.proto was not found in the cmvr-es toolchain"
)
def main() -> int:
args = parse_args()
proto_root = args.cmvr_es_root.resolve() / "protos"
if not proto_root.is_dir():
raise FileNotFoundError(f"proto root does not exist: {proto_root}")
proto_files = sorted((proto_root / "cmvr").rglob("*.proto"))
if not proto_files:
raise FileNotFoundError(f"no cmvr proto files found under {proto_root}")
args.output.mkdir(parents=True, exist_ok=True)
command = build_command(args, proto_files)
subprocess.run(command, cwd=proto_root, check=True)
print(f"generated {len(proto_files)} proto files into {args.output.resolve()}")
return 0
if __name__ == "__main__":
raise SystemExit(main())