[BUG](codegen) Unwrap RootModel fields to their bare root shape

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.

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),

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.
 

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<string,int>`), 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])
 
 

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.
 

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

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)