diff --git a/packages/overture-schema-codegen/docs/walkthrough.md b/packages/overture-schema-codegen/docs/walkthrough.md index 7cfabf84e..e5be7e6d0 100644 --- a/packages/overture-schema-codegen/docs/walkthrough.md +++ b/packages/overture-schema-codegen/docs/walkthrough.md @@ -144,6 +144,12 @@ outermost structural layer, which is exactly the `ArrayOf` that was just constru **dict** recurses separately for key and value types (with `newtype_ctx=None` for both, since dict keys and values are independent spines) and returns `MapOf`. +**RootModel** subclasses are handled distinctly from other `BaseModel` terminals, since a +RootModel serializes as its bare root value. `_unwrap` intercepts them just before the +terminal, recurses into the `root` field's annotation, and reattaches any root metadata with +`attach_field_metadata` -- exactly as a model field's own metadata reattaches. A +`RootModel[dict[str, int]]` field yields a bare `MapOf` carrying the root type's shape. + **Terminal** classification in `_terminal` handles the base case: `Any` becomes `AnyScalar`, `Literal` becomes `LiteralScalar`, `BaseModel` subclasses route through `model_resolver` (or fall back to `Primitive(source_type=cls)`), everything else becomes @@ -275,7 +281,7 @@ classes. One subtlety: Pydantic strips the `Annotated` wrapper from some fields and moves the metadata to `field_info.metadata`. When this happens, `analyze_type` sees a bare type -and misses the constraints. `_attach_field_metadata` routes them through +and misses the constraints. `attach_field_metadata` routes them through `attach_constraints` -- tagging them with `source=None` since they came from the field's own annotation rather than a NewType chain -- so length-constraint typing happens here just as it does during normal `Annotated` unwrapping. diff --git a/packages/overture-schema-codegen/src/overture/schema/codegen/extraction/model_extraction.py b/packages/overture-schema-codegen/src/overture/schema/codegen/extraction/model_extraction.py index 1f4a28167..ee8f47d00 100644 --- a/packages/overture-schema-codegen/src/overture/schema/codegen/extraction/model_extraction.py +++ b/packages/overture-schema-codegen/src/overture/schema/codegen/extraction/model_extraction.py @@ -12,8 +12,6 @@ from .docstring import clean_docstring from .field import ( - ConstraintSource, - FieldShape, ModelRef, UnionRef, ) @@ -22,7 +20,7 @@ ModelResolver, UnionResolver, analyze_type, - attach_constraints, + attach_field_metadata, unwrap_list, ) @@ -57,23 +55,6 @@ def _is_field_required(field_info: FieldInfo, is_optional: bool) -> bool: return not has_default and not is_optional -def _attach_field_metadata(shape: FieldShape, field_info: FieldInfo) -> FieldShape: - """Merge constraints from `field_info.metadata` onto *shape*. - - Pydantic strips the outermost Annotated wrapper from some fields - (non-optional, non-union) and moves its metadata to - `field_info.metadata`. When that happens `analyze_type` sees a bare - type and misses those constraints. They anchor at the topmost - constraint-bearing layer, so we route them through - `attach_constraints` so that length-constraint wrapping applies here - just as it does during normal annotation unwrapping. - """ - if not field_info.metadata: - return shape - extra = tuple(ConstraintSource(None, None, m) for m in field_info.metadata) - return attach_constraints(shape, extra) - - def _basemodel_bases(cls: type) -> list[type[BaseModel]]: """Return direct BaseModel bases, excluding BaseModel itself.""" return [b for b in cls.__bases__ if is_model_class(b) and b is not BaseModel] @@ -185,7 +166,12 @@ def _extract_model_recursive( model_resolver=model_resolver, union_resolver=union_resolver, ) - shape = _attach_field_metadata(shape, field_info) + # Pydantic strips the outermost Annotated wrapper from some fields + # (non-optional, non-union) and moves its metadata to + # `field_info.metadata`; `analyze_type` then sees a bare type and + # misses those constraints. Reattach them at the topmost + # constraint-bearing layer. + shape = attach_field_metadata(shape, field_info) fields.append( FieldSpec( name=resolve_field_alias(field_name, field_info), diff --git a/packages/overture-schema-codegen/src/overture/schema/codegen/extraction/type_analyzer.py b/packages/overture-schema-codegen/src/overture/schema/codegen/extraction/type_analyzer.py index 7c3c6a71d..559ac62c7 100644 --- a/packages/overture-schema-codegen/src/overture/schema/codegen/extraction/type_analyzer.py +++ b/packages/overture-schema-codegen/src/overture/schema/codegen/extraction/type_analyzer.py @@ -24,6 +24,10 @@ `UnsupportedUnionError`. Callers that need to recurse into sub-models pass resolvers that build a `ModelRef`/`UnionRef` with the resolved spec. + +A `RootModel` never reaches those terminals: it serializes as its bare +root value, so it is unwrapped to the root type's shape (with any root +metadata reattached) before terminal classification -- resolver or not. """ from __future__ import annotations @@ -43,7 +47,7 @@ ) from annotated_types import MaxLen, MinLen -from pydantic import BaseModel +from pydantic import BaseModel, RootModel from pydantic.fields import FieldInfo from typing_extensions import Sentinel, assert_never, evaluate_forward_ref @@ -103,6 +107,7 @@ class _NewTypeCtx: "UnsupportedUnionError", "analyze_type", "attach_constraints", + "attach_field_metadata", "capture_union_members", "is_newtype", "single_literal_value", @@ -371,6 +376,16 @@ def _recurse( value_shape, _, _ = _recurse(args[1], None) return MapOf(key=key_shape, value=value_shape, constraints=()), False, None + if isinstance(annotation, type) and issubclass(annotation, RootModel): + # A RootModel serializes as its bare root value, so unwrap to the + # root type's shape. Root-level constraints reattach exactly as + # field metadata does, so a constrained root + # (`RootModel[Annotated[list, MaxLen]]`) keeps its length-wrapped + # variant on the unwrapped layer. + root = annotation.model_fields["root"] + inner, opt, desc = _recurse(root.annotation, newtype_ctx) + return attach_field_metadata(inner, root), opt, desc + return _terminal(annotation, newtype_ctx, model_resolver), False, None @@ -492,6 +507,20 @@ def attach_constraints( assert_never(shape) +def attach_field_metadata(shape: FieldShape, field_info: FieldInfo) -> FieldShape: + """Merge constraints from `field_info.metadata` onto *shape*. + + Routes the metadata through `attach_constraints` so length-constraint + wrapping applies here just as it does during normal annotation + unwrapping: the constraints anchor at the topmost constraint-bearing + layer. Returns *shape* unchanged when there is no metadata. + """ + if not field_info.metadata: + return shape + extra = tuple(ConstraintSource(None, None, m) for m in field_info.metadata) + return attach_constraints(shape, extra) + + def _wrap_length_for_array(cs: ConstraintSource) -> ConstraintSource: """Replace a raw `MinLen`/`MaxLen` with its `ArrayOf`-layer variant. diff --git a/packages/overture-schema-codegen/tests/codegen_test_support.py b/packages/overture-schema-codegen/tests/codegen_test_support.py index c30bffebb..2050bb5d0 100644 --- a/packages/overture-schema-codegen/tests/codegen_test_support.py +++ b/packages/overture-schema-codegen/tests/codegen_test_support.py @@ -56,7 +56,7 @@ ) from overture.schema.system.ref import Id, Identified, Reference, Relationship from overture.schema.system.string import HexColor, LanguageTag, StrippedString -from pydantic import BaseModel, EmailStr, Field, HttpUrl +from pydantic import BaseModel, EmailStr, Field, HttpUrl, RootModel STR_TYPE = Primitive(base_type="str") @@ -199,6 +199,22 @@ class Widget(BaseModel): label: str = Field(description="Display label") +class TollChargesByVehicleType(RootModel[dict[str, int]]): + """A map-rooted RootModel: a bare `dict[str, int]` in serialized data.""" + + +class FeatureWithRootModel(BaseModel): + """A feature carrying a `RootModel`-typed field. + + The field serializes as its bare root value (a `map`), so + codegen must extract it to that shape rather than a struct with a + synthetic `root` member. + """ + + road_class: str + toll_charges: TollChargesByVehicleType | None = None + + CommonNames = NewType("CommonNames", dict[LanguageTag, StrippedString]) diff --git a/packages/overture-schema-codegen/tests/test_model_extraction.py b/packages/overture-schema-codegen/tests/test_model_extraction.py index 5ee081459..e9b23a3ba 100644 --- a/packages/overture-schema-codegen/tests/test_model_extraction.py +++ b/packages/overture-schema-codegen/tests/test_model_extraction.py @@ -2,8 +2,10 @@ from typing import Annotated, Optional +from codegen_test_support import FeatureWithRootModel from overture.schema.codegen.extraction.field import ( ArrayOf, + MapOf, ModelRef, Primitive, UnionRef, @@ -29,6 +31,21 @@ class Container(BaseModel): assert terminal.union.discriminator_field == "dimension" +def test_rootmodel_field_extracts_bare_root() -> None: + """A `RootModel`-typed field extracts to its bare root shape. + + Pydantic validates and serializes a RootModel as its bare root value, + so extraction must not produce a `ModelRef` struct with a synthetic + `root` member -- the generated schema would then declare a wrapper the + data never carries. + """ + spec = extract_model(FeatureWithRootModel) + toll = next(f for f in spec.fields if f.name == "toll_charges") + + assert isinstance(toll.shape, MapOf) + assert toll.is_optional is True + + def test_required_list_with_optional_element_is_required() -> None: """A required `list[X | None]` field must not inherit element optionality. diff --git a/packages/overture-schema-codegen/tests/test_pyspark_schema_builder.py b/packages/overture-schema-codegen/tests/test_pyspark_schema_builder.py index 825176c5d..58c96c897 100644 --- a/packages/overture-schema-codegen/tests/test_pyspark_schema_builder.py +++ b/packages/overture-schema-codegen/tests/test_pyspark_schema_builder.py @@ -3,7 +3,7 @@ from enum import Enum import pytest -from codegen_test_support import spec_for_model +from codegen_test_support import FeatureWithRootModel, spec_for_model from overture.schema.codegen.extraction.field import Primitive from overture.schema.codegen.extraction.specs import ( AnnotatedField, @@ -93,6 +93,23 @@ def test_dict_str_str_maps_to_map_type(self, fields: list[SchemaField]) -> None: assert labels_field.type_expr == "MapType(StringType(), StringType(), True)" +class TestRootModelField: + """A `RootModel` field renders as its bare root type, not a struct. + + Pydantic serializes a RootModel as its bare root value, so the Spark + schema must declare the value's type directly -- `MapType(...)` here, + never `StructType([StructField("root", ...)])`. + """ + + @pytest.fixture + def fields(self) -> list[SchemaField]: + return build_schema(spec_for_model(FeatureWithRootModel)) + + def test_rootmodel_field_maps_to_map_type(self, fields: list[SchemaField]) -> None: + toll = next(f for f in fields if f.name == "toll_charges") + assert toll.type_expr == "MapType(StringType(), LongType(), True)" + + class TestDivisionAreaSchema: @pytest.fixture(scope="class") @classmethod diff --git a/packages/overture-schema-codegen/tests/test_type_analyzer.py b/packages/overture-schema-codegen/tests/test_type_analyzer.py index 144214b75..869e2d000 100644 --- a/packages/overture-schema-codegen/tests/test_type_analyzer.py +++ b/packages/overture-schema-codegen/tests/test_type_analyzer.py @@ -5,6 +5,7 @@ import pytest from annotated_types import Ge, MaxLen, MinLen +from codegen_test_support import TollChargesByVehicleType from overture.schema.codegen.extraction.field import ( AnyScalar, ArrayOf, @@ -22,6 +23,7 @@ list_depth, ) from overture.schema.codegen.extraction.length_constraints import ( + ArrayMaxLen, ArrayMinLen, ScalarMinLen, ) @@ -46,7 +48,7 @@ NoWhitespaceString, SnakeCaseString, ) -from pydantic import BaseModel, Field, Tag +from pydantic import BaseModel, Field, RootModel, Tag from typing_extensions import Sentinel @@ -285,6 +287,68 @@ class Person(BaseModel): assert shape.base_type == "Person" +class TestRootModel: + """A `RootModel` serializes as its bare root value. + + So `analyze_type` unwraps it to the root type's shape, reattaching any + root-level metadata onto the unwrapped layer. `TollChargesByVehicleType` + is a map-rooted RootModel; the scalar and constrained roots below are + local because each exercises one path only. + """ + + def test_map_root_unwraps_to_mapof(self) -> None: + shape = _shape(TollChargesByVehicleType) + assert isinstance(shape, MapOf) + assert isinstance(shape.key, Primitive) and shape.key.base_type == "str" + assert isinstance(shape.value, Primitive) and shape.value.base_type == "int" + + def test_scalar_root_unwraps_to_primitive(self) -> None: + class Slug(RootModel[str]): + pass + + shape = _shape(Slug) + assert isinstance(shape, Primitive) + assert shape.base_type == "str" + + def test_constrained_root_reattaches_constraints(self) -> None: + class Tags(RootModel[Annotated[list[str], MaxLen(3)]]): + pass + + shape = _shape(Tags) + assert isinstance(shape, ArrayOf) + assert ArrayMaxLen in {type(cs.constraint) for cs in shape.constraints} + + def test_constrained_scalar_root_reattaches_scalar_min_len(self) -> None: + class Code(RootModel[Annotated[str, MinLen(2)]]): + pass + + shape = _shape(Code) + assert isinstance(shape, Primitive) + assert ScalarMinLen in {type(cs.constraint) for cs in shape.constraints} + + def test_unwrap_bypasses_model_resolver(self) -> None: + """The RootModel is unwrapped before terminal classification. + + A plain `BaseModel` terminal routes through `model_resolver`; a + RootModel resolves structurally into its root shape, so the + resolver is never invoked for it. + """ + seen: list[type] = [] + + def resolver(cls: type[BaseModel]) -> FieldShape: + seen.append(cls) + return ModelRef(model=RecordSpec(name=cls.__name__, description=None)) + + shape, _, _ = analyze_type(TollChargesByVehicleType, model_resolver=resolver) + assert isinstance(shape, MapOf) + assert seen == [] + + def test_optional_rootmodel_field_is_optional(self) -> None: + shape, is_optional, _ = analyze_type(TollChargesByVehicleType | None) + assert isinstance(shape, MapOf) + assert is_optional is True + + class TestNewType: def test_simple_newtype(self) -> None: shape = _shape(int32)