fix(parsing): parse extra field types

Author: stainless-app[bot]

src/writerai/_models.py

@@ -209,14 +209,18 @@ def construct(  # pyright: ignore[reportIncompatibleMethodOverride]
             else:
                 fields_values[name] = field_get_default(field)
 
+        extra_field_type = _get_extra_fields_type(__cls)
+
         _extra = {}
         for key, value in values.items():
             if key not in model_fields:
+                parsed = construct_type(value=value, type_=extra_field_type) if extra_field_type is not None else value
+
                 if PYDANTIC_V2:
-                    _extra[key] = value
+                    _extra[key] = parsed
                 else:
                     _fields_set.add(key)
-                    fields_values[key] = value
+                    fields_values[key] = parsed
 
         object.__setattr__(m, "__dict__", fields_values)
 
@@ -371,6 +375,23 @@ def _construct_field(value: object, field: FieldInfo, key: str) -> object:
     return construct_type(value=value, type_=type_, metadata=getattr(field, "metadata", None))
 
 
+def _get_extra_fields_type(cls: type[pydantic.BaseModel]) -> type | None:
+    if not PYDANTIC_V2:
+        # TODO
+        return None
+
+    schema = cls.__pydantic_core_schema__
+    if schema["type"] == "model":
+        fields = schema["schema"]
+        if fields["type"] == "model-fields":
+            extras = fields.get("extras_schema")
+            if extras and "cls" in extras:
+                # mypy can't narrow the type
+                return extras["cls"]  # type: ignore[no-any-return]
+
+    return None
+
+
 def is_basemodel(type_: type) -> bool:
     """Returns whether or not the given type is either a `BaseModel` or a union of `BaseModel`"""
     if is_union(type_):

tests/test_models.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict, List, Union, Optional, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Union, Optional, cast
 from datetime import datetime, timezone
 from typing_extensions import Literal, Annotated, TypeAliasType
 
@@ -934,3 +934,30 @@ class Type2(BaseModel):
     )
     assert isinstance(model, Type1)
     assert isinstance(model.value, InnerType2)
+
+
+@pytest.mark.skipif(not PYDANTIC_V2, reason="this is only supported in pydantic v2 for now")
+def test_extra_properties() -> None:
+    class Item(BaseModel):
+        prop: int
+
+    class Model(BaseModel):
+        __pydantic_extra__: Dict[str, Item] = Field(init=False)  # pyright: ignore[reportIncompatibleVariableOverride]
+
+        other: str
+
+        if TYPE_CHECKING:
+
+            def __getattr__(self, attr: str) -> Item: ...
+
+    model = construct_type(
+        type_=Model,
+        value={
+            "a": {"prop": 1},
+            "other": "foo",
+        },
+    )
+    assert isinstance(model, Model)
+    assert model.a.prop == 1
+    assert isinstance(model.a, Item)
+    assert model.other == "foo"