mizan/protocol/mizan-codegen/python/pydantic_to_rust/emit.py

"""`ModelDecl` -> Rust source text.

This module classifies shapes and defaults; `templates/` spells the Rust.
Emitted structs are for serialization — plain serde derives, no `repr(C)`,
no padding, no alignment control.
"""

from __future__ import annotations

from pathlib import Path

from jinja2 import Environment, FileSystemLoader, StrictUndefined

from pydantic_to_rust.idents import to_rust_ident
from pydantic_to_rust.ir import DefaultValue, ModelDecl, TypeShape

_DEFAULT_DERIVES = (
    "Debug",
    "Clone",
    "::serde::Serialize",
    "::serde::Deserialize",
)

_TEMPLATES = Path(__file__).parent / "templates"


def _rust_string(value: str) -> str:
    escaped = value.replace("\\", "\\\\").replace('"', '\\"')
    return f'"{escaped}"'


def _rust_f64(value: float) -> str:
    """Always carry a decimal point, so `0.65` cannot read as an integer."""
    if value == int(value):
        return f"{int(value)}.0"
    return repr(float(value))


def _environment() -> Environment:
    env = Environment(
        loader=FileSystemLoader(_TEMPLATES),
        undefined=StrictUndefined,
        keep_trailing_newline=True,
    )
    env.filters["rust_string"] = _rust_string
    env.filters["rust_f64"] = _rust_f64
    return env


_ENV = _environment()


def _classify_shape(shape: TypeShape) -> dict:
    """One node per reachable `TypeShape` variant, for the type template."""
    if shape.primitive is not None:
        return {"kind": "primitive", "primitive": shape.primitive.value}
    if shape.ref is not None:
        return {"kind": "ref", "name": shape.ref}
    if shape.array_element is not None and shape.array_length is not None:
        return {
            "kind": "array",
            "element": _classify_shape(shape.array_element),
            "length": shape.array_length,
        }
    if shape.list_inner is not None:
        return {"kind": "vec", "inner": _classify_shape(shape.list_inner)}
    if shape.optional_inner is not None:
        return {"kind": "option", "inner": _classify_shape(shape.optional_inner)}
    if shape.enum_variants is not None:
        # An inline `Literal[...]` has no name to emit.
        return {"kind": "primitive", "primitive": "string"}
    if shape.union_branches is not None:
        return {"kind": "opaque"}
    if shape.map_value is not None:
        key = (
            _classify_shape(shape.map_key)
            if shape.map_key is not None
            else {"kind": "primitive", "primitive": "string"}
        )
        return {"kind": "map", "key": key, "value": _classify_shape(shape.map_value)}
    raise ValueError(f"no variant set on TypeShape: {shape}")


def _classify_value(value) -> dict:
    """A JSON-shaped value, for the value template."""
    if value is None:
        return {"kind": "unit"}
    if isinstance(value, bool):
        return {"kind": "bool", "value": value}
    if isinstance(value, int):
        return {"kind": "int", "value": value}
    if isinstance(value, float):
        return {"kind": "float", "value": value}
    if isinstance(value, str):
        return {"kind": "string", "value": value}
    if isinstance(value, (list, tuple)):
        if not value:
            return {"kind": "empty_vec"}
        return {"kind": "vec", "elements": [_classify_value(v) for v in value]}
    if isinstance(value, dict):
        if not value:
            return {"kind": "empty_map"}
        return {
            "kind": "map",
            "pairs": [
                {"key": k, "value": _classify_value(v)} for k, v in value.items()
            ],
        }
    return {"kind": "unit"}


def _classify_default(shape: TypeShape, default: DefaultValue) -> dict:
    """One node per default kind, for the default template."""
    if default.kind == "null":
        if shape.optional_inner is not None:
            return {"kind": "none"}
        return {"kind": "unit"}
    if default.kind == "boolean":
        return {"kind": "bool", "value": default.literal}
    if default.kind == "integer":
        return {"kind": "int", "value": int(default.literal)}
    if default.kind == "number":
        return {"kind": "float", "value": default.literal}
    if default.kind == "string":
        return {"kind": "string", "value": default.literal}
    if default.kind == "empty_seq":
        return {"kind": "empty_vec"}
    if default.kind == "empty_map":
        return {"kind": "empty_map"}
    if default.kind == "tuple":
        return {
            "kind": "array",
            "elements": [_classify_value(v) for v in default.literal],
        }
    if default.kind == "enum_variant":
        return {
            "kind": "enum_variant",
            "type": default.literal["type"],
            "variant": default.literal["variant"],
        }
    if default.kind == "compound_model":
        # Defer to the referenced type's own Default rather than replaying
        # the factory's field values.
        return {"kind": "delegate", "type": default.literal["type"]}
    return {"kind": "unit"}


def emit_rust_struct(
    model: ModelDecl,
    *,
    derives: tuple[str, ...] = _DEFAULT_DERIVES,
    extra_attrs: tuple[str, ...] = (),
) -> str:
    """Render one model as a struct, its default functions, and an `impl
    Default`. Returns a single block; the caller joins blocks into a file."""
    fields = []
    default_fns = []
    for decl in model.fields:
        shape = _classify_shape(decl.shape)
        default_fn = None
        if decl.default is not None:
            default_fn = f"__default_{model.name}_{decl.name}"
            default_fns.append(
                {
                    "name": default_fn,
                    "type": shape,
                    "body": _classify_default(decl.shape, decl.default),
                }
            )
        fields.append(
            {
                "ident": to_rust_ident(decl.name),
                "type": shape,
                "default_fn": default_fn,
            }
        )

    doc_lines = []
    if model.docstring:
        doc_lines = [line.rstrip() for line in model.docstring.strip().splitlines()]

    return _ENV.get_template("struct.rs.j2").render(
        name=model.name,
        doc_lines=doc_lines,
        derives=list(derives),
        extra_attrs=list(extra_attrs),
        fields=fields,
        default_fns=default_fns,
    )