Handle __pydantic_extra__ annotation being a string or inherited

pydantic · Jan 29, 2024 · edad102 · edad102
1 parent b785d5b
commit edad102
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/pydantic/_internal/_generate_schema.py b/pydantic/_internal/_generate_schema.py
@@ -538,16 +538,23 @@ def _model_schema(self, cls: type[BaseModel]) -> core_schema.CoreSchema:
 
             extras_schema = None
             if core_config.get('extra_fields_behavior') == 'allow':
-                for tp in (cls, *cls.__mro__):
-                    extras_annotation = cls.__annotations__.get('__pydantic_extra__', None)
+                assert cls.__mro__[0] is cls
+                assert cls.__mro__[-1] is object
+                for candidate_cls in cls.__mro__[:-1]:
+                    extras_annotation = candidate_cls.__annotations__.get('__pydantic_extra__', None)
                     if extras_annotation is not None:
+                        if isinstance(extras_annotation, str):
+                            extras_annotation = _typing_extra.eval_type_backport(
+                                _typing_extra._make_forward_ref(extras_annotation, is_argument=False, is_class=True),
+                                self._types_namespace,
+                            )
                         tp = get_origin(extras_annotation)
                         if tp not in (Dict, dict):
                             raise PydanticSchemaGenerationError(
                                 'The type annotation for `__pydantic_extra__` must be `Dict[str, ...]`'
                             )
                         extra_items_type = self._get_args_resolving_forward_refs(
-                            cls.__annotations__['__pydantic_extra__'],
+                            extras_annotation,
                             required=True,
                         )[1]
                         if extra_items_type is not Any:

diff --git a/tests/test_main.py b/tests/test_main.py
@@ -2842,11 +2842,12 @@ class Foo(BaseModel):
 
     class Model(BaseModel):
         model_config = ConfigDict(extra='allow')
+        __pydantic_extra__: 'dict[str, Foo]'
 
     class Child(Model):
-        __pydantic_extra__: Dict[str, Foo]
+        y: int
 
-    m = Child(a={'x': '1'})
+    m = Child(a={'x': '1'}, y=2)
     assert m.__pydantic_extra__ == {'a': Foo(x=1)}
 
     # insert_assert(Child.model_json_schema())
@@ -2860,7 +2861,8 @@ class Child(Model):
             }
         },
         'additionalProperties': {'$ref': '#/$defs/Foo'},
-        'properties': {},
+        'properties': {'y': {'title': 'Y', 'type': 'integer'}},
+        'required': ['y'],
         'title': 'Child',
         'type': 'object',
     }