Files
mizan/protocol/mizan-codegen/scripts/run_decoru.py
2026-06-04 01:15:41 -04:00

140 lines
4.8 KiB
Python

#!/usr/bin/env python3
"""Pydantic → Rust codegen helper invoked by mizan-codegen's
`[source.rust.pydantic]` step.
Reads a JSON payload from argv[1] with keys:
- module: Python module to import (e.g. "claude_manage.schema")
- output: Path to write the generated Rust file
- derives: List of derive identifiers to apply to every emitted item
- header: Optional file prefix (e.g. an AUTO-GENERATED warning)
Discovers every BaseModel subclass declared in the module (handled by
decoru) AND every Enum subclass declared there (handled inline — decoru
itself is scoped to BaseModel). Writes one Rust file containing both.
Bundled with the mizan-codegen binary (include_str!) and piped to
`python -` at codegen time — no install step beyond decoru being
importable in the python environment.
"""
import importlib
import inspect
import json
import sys
import textwrap
from enum import Enum
from pathlib import Path
from pydantic import BaseModel # type: ignore[import-untyped]
from decoru import ( # type: ignore[import-untyped]
emit_rust_struct,
to_rust_variant_ident,
walk_pydantic_model,
)
def _declared_in(module, obj) -> bool:
return getattr(obj, "__module__", None) == module.__name__
def discover_models(module) -> list[type[BaseModel]]:
"""BaseModel subclasses declared in this module. Imported helpers are
skipped — only own-module declarations qualify."""
return [
obj
for _, obj in inspect.getmembers(module, inspect.isclass)
if issubclass(obj, BaseModel)
and obj is not BaseModel
and _declared_in(module, obj)
]
def discover_enums(module) -> list[type[Enum]]:
"""Enum subclasses declared in this module. Filters out the
framework's own Enum class and anything imported from elsewhere."""
return [
obj
for _, obj in inspect.getmembers(module, inspect.isclass)
if issubclass(obj, Enum) and obj is not Enum and _declared_in(module, obj)
]
# Last-variant-is-default matches the catch-all idiom (e.g. `Metadata`
# in `claude_manage.schema.EntryType`). Decoru's `emit_rust_struct`
# emits `impl Default` unconditionally on every BaseModel, so any
# enum-typed field that lacks a Pydantic default must still satisfy
# `EntryType::default()`. Forcing #[default] on the last member keeps
# the generated structs compilable without per-enum config.
_ENUM_TEMPLATE = textwrap.dedent("""\
#[derive({derives})]
#[serde(rename_all = "snake_case")]
pub enum {name} {{
{variants}
}}
""")
def _render_variant(member: Enum, *, is_default: bool) -> str:
# Pascal-casing the Python member name is the same conversion decoru
# applies when capturing enum field defaults. Sharing the function is
# load-bearing — divergent conversions emit non-compiling schema.rs.
pascal = to_rust_variant_ident(member.name)
default_attr = " #[default]\n" if is_default else ""
return f"{default_attr} {pascal},"
def emit_rust_enum(enum_class: type[Enum], derives: tuple[str, ...]) -> str:
"""Render a Rust enum with PascalCase variants from Python member
names. Pairs `#[serde(rename_all = "snake_case")]` so the wire form
matches each member's `value`. Adds `Default` to the derives and
marks the last member `#[default]` — see `_ENUM_TEMPLATE` for the
rationale."""
name = enum_class.__name__
members = list(enum_class)
full_derives = ", ".join((*derives, "Default"))
variants = "\n".join(
_render_variant(m, is_default=(i == len(members) - 1))
for i, m in enumerate(members)
)
return _ENUM_TEMPLATE.format(derives=full_derives, name=name, variants=variants)
def main() -> int:
if len(sys.argv) < 2:
sys.stderr.write("run_decoru.py: missing JSON payload argument\n")
return 2
payload = json.loads(sys.argv[1])
module_name: str = payload["module"]
output_path = Path(payload["output"]).resolve()
derives = tuple(payload.get("derives", ()))
header = payload.get("header") or ""
sys.path.insert(0, str(Path.cwd()))
module = importlib.import_module(module_name)
enums = discover_enums(module)
models = discover_models(module)
if not enums and not models:
sys.stderr.write(
f"run_decoru.py: no Enum or BaseModel subclasses declared in {module_name!r}\n"
)
return 3
enum_blocks = [emit_rust_enum(e, derives) for e in enums]
struct_blocks = [emit_rust_struct(walk_pydantic_model(m), derives=derives) for m in models]
body = "\n".join((*enum_blocks, *struct_blocks))
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(header + body)
sys.stderr.write(
f"run_decoru.py: wrote {len(enums)} enum(s) + {len(models)} struct(s) to {output_path}\n"
)
return 0
if __name__ == "__main__":
sys.exit(main())