Fixes/upgrades
This commit is contained in:
285
protocol/mizan-codegen/python/pydantic_to_rust/walker.py
Normal file
285
protocol/mizan-codegen/python/pydantic_to_rust/walker.py
Normal file
@@ -0,0 +1,285 @@
|
||||
"""Pydantic class -> `ModelDecl`.
|
||||
|
||||
Walks `model_fields`, lowers each annotation to a `TypeShape`, and captures
|
||||
defaults. Annotations this does not recognize lower to `Primitive.STRING`,
|
||||
matching what Mizan's own IR does with an unknown shape.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
import re
|
||||
import sys
|
||||
import types
|
||||
from typing import Any, Literal, Union, get_args, get_origin
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic_core import PydanticUndefined
|
||||
|
||||
from pydantic_to_rust.idents import to_rust_variant_ident
|
||||
from pydantic_to_rust.ir import (
|
||||
DefaultValue,
|
||||
FieldDecl,
|
||||
ModelDecl,
|
||||
Primitive,
|
||||
TypeShape,
|
||||
)
|
||||
|
||||
# NewTypes named i8..i64 / u8..u64 / f16..f64, recognized only when the
|
||||
# caller opts in.
|
||||
_WIDTH_NEWTYPE_RE = re.compile(r"^([iuf])(8|16|32|64)$")
|
||||
|
||||
|
||||
def walk_pydantic_model(
|
||||
model: type[BaseModel],
|
||||
*,
|
||||
newtype_width_detection: bool = False,
|
||||
) -> ModelDecl:
|
||||
"""Lower one Pydantic class.
|
||||
|
||||
With `newtype_width_detection`, width-named NewTypes carry their width
|
||||
and signedness. Off, every NewType lowers to a `ref`.
|
||||
"""
|
||||
fields: list[FieldDecl] = []
|
||||
for field_name, field_info in model.model_fields.items():
|
||||
required = field_info.is_required()
|
||||
default: DefaultValue | None = None
|
||||
if not required:
|
||||
# A factory wins over a literal — it is the canonical Pydantic
|
||||
# path for compound defaults.
|
||||
factory = getattr(field_info, "default_factory", None)
|
||||
if factory is not None:
|
||||
default = _capture_factory_default(_call_factory(factory, field_name))
|
||||
else:
|
||||
default = _capture_default(field_info.default)
|
||||
fields.append(
|
||||
FieldDecl(
|
||||
name=field_name,
|
||||
shape=lower_annotation(
|
||||
field_info.annotation,
|
||||
newtype_width_detection=newtype_width_detection,
|
||||
),
|
||||
required=required,
|
||||
default=default,
|
||||
)
|
||||
)
|
||||
return ModelDecl(
|
||||
name=model.__name__,
|
||||
fields=tuple(fields),
|
||||
docstring=(model.__doc__ or None),
|
||||
)
|
||||
|
||||
|
||||
def _call_factory(factory, field_name: str):
|
||||
"""A factory that raises leaves the field with no captured default."""
|
||||
try:
|
||||
return factory()
|
||||
except Exception as exc:
|
||||
print(
|
||||
f"pydantic_to_rust: default_factory for {field_name!r} raised "
|
||||
f"{type(exc).__name__}: {exc}; emitting no default",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def lower_annotation(
|
||||
annotation: Any,
|
||||
*,
|
||||
newtype_width_detection: bool = False,
|
||||
) -> TypeShape:
|
||||
"""Lower a bare annotation, for callers that do not have a whole model."""
|
||||
return _lower_annotation(
|
||||
annotation,
|
||||
newtype_width_detection=newtype_width_detection,
|
||||
)
|
||||
|
||||
|
||||
def _capture_factory_default(value: Any) -> DefaultValue | None:
|
||||
"""Capture a factory's result. Empty collections map onto Rust's own
|
||||
`Default`; a model instance is serialized so the emitter can name it."""
|
||||
if value is None:
|
||||
return DefaultValue(kind="null", literal=None)
|
||||
if isinstance(value, (list, tuple, set, frozenset)) and len(value) == 0:
|
||||
return DefaultValue(kind="empty_seq", literal=None)
|
||||
if isinstance(value, dict) and len(value) == 0:
|
||||
return DefaultValue(kind="empty_map", literal=None)
|
||||
if isinstance(value, BaseModel):
|
||||
return DefaultValue(
|
||||
kind="compound_model",
|
||||
literal={"type": type(value).__name__, "fields": value.model_dump()},
|
||||
)
|
||||
if isinstance(value, bool):
|
||||
return DefaultValue(kind="boolean", literal=value)
|
||||
if isinstance(value, int):
|
||||
return DefaultValue(kind="integer", literal=value)
|
||||
if isinstance(value, float):
|
||||
return DefaultValue(kind="number", literal=value)
|
||||
if isinstance(value, str):
|
||||
return DefaultValue(kind="string", literal=value)
|
||||
return DefaultValue(kind="compound", literal=value)
|
||||
|
||||
|
||||
def _capture_default(raw: Any) -> DefaultValue | None:
|
||||
"""Filter Pydantic's sentinels, then record the literal's kind."""
|
||||
if raw is None:
|
||||
return DefaultValue(kind="null", literal=None)
|
||||
if raw is PydanticUndefined or raw is ...:
|
||||
return None
|
||||
if isinstance(raw, enum.Enum):
|
||||
# Store the Rust variant name, not the Python member name, so the
|
||||
# default body and the enum declaration agree.
|
||||
return DefaultValue(
|
||||
kind="enum_variant",
|
||||
literal={
|
||||
"type": type(raw).__name__,
|
||||
"variant": to_rust_variant_ident(raw.name),
|
||||
},
|
||||
)
|
||||
if isinstance(raw, bool):
|
||||
return DefaultValue(kind="boolean", literal=raw)
|
||||
if isinstance(raw, int):
|
||||
return DefaultValue(kind="integer", literal=raw)
|
||||
if isinstance(raw, float):
|
||||
return DefaultValue(kind="number", literal=raw)
|
||||
if isinstance(raw, str):
|
||||
return DefaultValue(kind="string", literal=raw)
|
||||
if isinstance(raw, tuple):
|
||||
return DefaultValue(kind="tuple", literal=tuple(raw))
|
||||
return DefaultValue(kind="compound", literal=raw)
|
||||
|
||||
|
||||
def _lower_annotation(
|
||||
annotation: Any,
|
||||
*,
|
||||
newtype_width_detection: bool = False,
|
||||
) -> TypeShape:
|
||||
# NewType carries `__supertype__` and `__name__`, and keeps its named
|
||||
# identity on the Rust side.
|
||||
if hasattr(annotation, "__supertype__") and hasattr(annotation, "__name__"):
|
||||
if newtype_width_detection:
|
||||
match = _WIDTH_NEWTYPE_RE.match(annotation.__name__)
|
||||
if match:
|
||||
kind, bits = match.group(1), int(match.group(2))
|
||||
if kind == "f":
|
||||
return TypeShape(primitive=Primitive.NUMBER, float_width=bits)
|
||||
return TypeShape(
|
||||
primitive=Primitive.INTEGER,
|
||||
int_width=bits,
|
||||
int_signed=(kind == "i"),
|
||||
)
|
||||
return TypeShape(ref=annotation.__name__)
|
||||
|
||||
inner, is_optional = _extract_optional(annotation)
|
||||
if is_optional:
|
||||
return TypeShape(
|
||||
optional_inner=_lower_annotation(
|
||||
inner, newtype_width_detection=newtype_width_detection
|
||||
)
|
||||
)
|
||||
|
||||
origin = get_origin(annotation)
|
||||
|
||||
if origin is Union or isinstance(annotation, types.UnionType):
|
||||
branches = [a for a in get_args(annotation) if a is not type(None)]
|
||||
if len(branches) > 1:
|
||||
return TypeShape(
|
||||
union_branches=tuple(
|
||||
_lower_annotation(
|
||||
b, newtype_width_detection=newtype_width_detection
|
||||
)
|
||||
for b in branches
|
||||
)
|
||||
)
|
||||
|
||||
# `tuple[T1..TN]` with no ellipsis becomes `[T; N]`, but only when every
|
||||
# element shares a shape; mixed tuples fall through to the list path.
|
||||
if origin is tuple:
|
||||
args = get_args(annotation)
|
||||
if args and args[-1] is not Ellipsis:
|
||||
first = _lower_annotation(
|
||||
args[0], newtype_width_detection=newtype_width_detection
|
||||
)
|
||||
uniform = all(
|
||||
_lower_annotation(a, newtype_width_detection=newtype_width_detection)
|
||||
== first
|
||||
for a in args[1:]
|
||||
)
|
||||
if uniform:
|
||||
return TypeShape(array_element=first, array_length=len(args))
|
||||
|
||||
element = _extract_list_element(annotation)
|
||||
if element is not None:
|
||||
return TypeShape(
|
||||
list_inner=_lower_annotation(
|
||||
element, newtype_width_detection=newtype_width_detection
|
||||
)
|
||||
)
|
||||
|
||||
if origin is dict:
|
||||
args = get_args(annotation)
|
||||
if len(args) == 2:
|
||||
key = _lower_annotation(
|
||||
args[0], newtype_width_detection=newtype_width_detection
|
||||
)
|
||||
value = _lower_annotation(
|
||||
args[1], newtype_width_detection=newtype_width_detection
|
||||
)
|
||||
return TypeShape(
|
||||
map_value=value,
|
||||
map_key=key if key.ref is not None else None,
|
||||
)
|
||||
|
||||
if origin is Literal:
|
||||
args = get_args(annotation)
|
||||
if all(isinstance(a, str) for a in args):
|
||||
return TypeShape(enum_variants=tuple(args))
|
||||
if len(args) == 1:
|
||||
sole = args[0]
|
||||
if isinstance(sole, bool):
|
||||
return TypeShape(primitive=Primitive.BOOLEAN)
|
||||
if isinstance(sole, int):
|
||||
return TypeShape(primitive=Primitive.INTEGER)
|
||||
if isinstance(sole, float):
|
||||
return TypeShape(primitive=Primitive.NUMBER)
|
||||
|
||||
# An Enum or model class is referenced by name; the caller emits or
|
||||
# hand-writes the corresponding Rust declaration.
|
||||
if isinstance(annotation, type) and issubclass(annotation, enum.Enum):
|
||||
return TypeShape(ref=annotation.__name__)
|
||||
if isinstance(annotation, type) and issubclass(annotation, BaseModel):
|
||||
return TypeShape(ref=annotation.__name__)
|
||||
|
||||
if annotation is int:
|
||||
return TypeShape(primitive=Primitive.INTEGER)
|
||||
if annotation is float:
|
||||
return TypeShape(primitive=Primitive.NUMBER)
|
||||
if annotation is bool:
|
||||
return TypeShape(primitive=Primitive.BOOLEAN)
|
||||
if annotation is str:
|
||||
return TypeShape(primitive=Primitive.STRING)
|
||||
|
||||
return TypeShape(primitive=Primitive.STRING)
|
||||
|
||||
|
||||
def _extract_optional(ann: Any) -> tuple[Any, bool]:
|
||||
"""`(T, True)` for `Optional[T]` or `T | None`, else `(ann, False)`."""
|
||||
origin = get_origin(ann)
|
||||
if origin is Union or isinstance(ann, types.UnionType):
|
||||
args = get_args(ann)
|
||||
non_none = [a for a in args if a is not type(None)]
|
||||
if len(non_none) == 1 and type(None) in args:
|
||||
return non_none[0], True
|
||||
return ann, False
|
||||
|
||||
|
||||
def _extract_list_element(ann: Any) -> Any | None:
|
||||
"""`T` for `list[T]`, `tuple[T, ...]`, `set[T]`, `frozenset[T]`."""
|
||||
origin = get_origin(ann)
|
||||
if origin in (list, tuple, set, frozenset):
|
||||
args = get_args(ann)
|
||||
if origin is tuple and len(args) >= 2 and args[1] is Ellipsis:
|
||||
return args[0]
|
||||
if args:
|
||||
return args[0]
|
||||
return None
|
||||
Reference in New Issue
Block a user