recursive models working, more tests fixed

pydantic · samuelcolvin · Nov 2, 2022 · Sep 12, 2022 · Sep 14, 2022 · Sep 14, 2022
commit 99a8a15f6b5698c3e99f807b91d8fe1be96378e5
diff --git a/pydantic/_internal/babelfish.py b/pydantic/_internal/babelfish.py
@@ -7,16 +7,7 @@
 import typing
 from typing import TYPE_CHECKING, Any
 
-from pydantic_core import Schema as PydanticCoreSchema
-from pydantic_core._types import (
-    Config as PydanticCoreConfig,
-    DictSchema,
-    IsInstanceSchema,
-    TuplePositionalSchema,
-    TupleVariableSchema,
-    TypedDictField,
-    TypedDictSchema,
-)
+from pydantic_core import schema_types as core_schema
 from typing_extensions import get_args, is_typeddict
 
 from pydantic.fields import FieldInfo, Undefined
@@ -40,25 +31,28 @@
 __all__ = 'generate_config', 'generate_schema', 'model_fields_schema'
 
 
-def model_fields_schema(fields: dict[str, FieldInfo], validator_functions: ValidationFunctions) -> PydanticCoreSchema:
-    schema: PydanticCoreSchema = {
-        'type': 'typed-dict',
-        'return_fields_set': True,
-        'fields': {k: generate_field_schema(k, v, validator_functions) for k, v in fields.items()},
-    }
+def model_fields_schema(
+    ref: str, fields: dict[str, FieldInfo], validator_functions: ValidationFunctions
+) -> core_schema.Schema:
+    schema: core_schema.Schema = core_schema.TypedDictSchema(
+        type='typed-dict',
+        ref=ref,
+        return_fields_set=True,
+        fields={k: generate_field_schema(k, v, validator_functions) for k, v in fields.items()},
+    )
     schema = apply_validators(schema, validator_functions.get_root_validators())
     return schema
 
 
-def generate_config(config: type[BaseConfig]) -> PydanticCoreConfig:
-    return {
-        'typed_dict_extra_behavior': config.extra.value,
-        # 'allow_inf_nan': config.allow_inf_nan,
-        'populate_by_name': config.allow_population_by_field_name,
-    }
+def generate_config(config: type[BaseConfig]) -> core_schema.Config:
+    return core_schema.Config(
+        typed_dict_extra_behavior=config.extra.value,
+        # allow_inf_nan=config.allow_inf_nan,
+        populate_by_name=config.allow_population_by_field_name,
+    )
 
 
-def generate_schema(obj: Any) -> PydanticCoreSchema:  # noqa: C901 (ignore complexity)
+def generate_schema(obj: Any) -> core_schema.Schema:  # noqa: C901 (ignore complexity)
     """
     Recursively generate a pydantic-core schema for any supported python type.
     """
@@ -72,7 +66,7 @@ def generate_schema(obj: Any) -> PydanticCoreSchema:  # noqa: C901 (ignore compl
     elif obj is None or obj is NoneType:
         return 'none'
     elif obj == type:
-        return {'type': 'is-instance', 'class_': type}
+        return core_schema.IsInstanceSchema(type='is-instance', class_=type)
     elif is_callable_type(obj):
         return 'callable'
     elif is_literal_type(obj):
@@ -90,6 +84,10 @@ def generate_schema(obj: Any) -> PydanticCoreSchema:  # noqa: C901 (ignore compl
     if schema_property is not None:
         return schema_property
 
+    get_schema = getattr(obj, '__get_pydantic_validation_schema__', None)
+    if get_schema is not None:
+        return get_schema()
+
     origin = get_origin(obj)
     if origin is None:
         raise PydanticSchemaGenerationError(f'Unable to generate pydantic-core schema for {obj!r}.')
@@ -108,62 +106,81 @@ def generate_schema(obj: Any) -> PydanticCoreSchema:  # noqa: C901 (ignore compl
         raise PydanticSchemaGenerationError(f'Unable to generate pydantic-core schema for {obj!r} (origin={origin!r}).')
 
 
-def generate_field_schema(name: str, field: FieldInfo, validator_functions: ValidationFunctions) -> TypedDictField:
+def generate_field_schema(
+    name: str, field: FieldInfo, validator_functions: ValidationFunctions
+) -> core_schema.TypedDictField:
     """
     Prepare a TypedDictField to represent a model or typeddict field.
     """
-    assert field.annotation is not None, 'field.annotation is None'
-    schema: PydanticCoreSchema = generate_schema(field.annotation)
+    assert field.annotation is not None, 'field.annotation should not be None when generating a schema'
+    schema: core_schema.Schema = generate_schema(field.annotation)
     schema = apply_validators(schema, validator_functions.get_field_validators(name))
 
-    field_schema: TypedDictField = {'schema': schema}
+    required = False
     if field.default_factory:
-        field_schema['default_factory'] = field.default_factory
+        schema = core_schema.WithDefaultSchema(type='default', schema=schema, default_factory=field.default_factory)
     elif field.default is Undefined:
-        field_schema['required'] = True
+        required = True
     else:
-        field_schema['default'] = field.default
+        schema = core_schema.WithDefaultSchema(type='default', schema=schema, default=field.default)
 
+    field_schema = core_schema.TypedDictField(schema=schema, required=required)
     if field.alias is not None:
         field_schema['alias'] = field.alias
     return field_schema
 
 
-def apply_validators(schema: PydanticCoreSchema, validators: list[Validator]) -> PydanticCoreSchema:
+def apply_validators(schema: core_schema.Schema, validators: list[Validator]) -> core_schema.Schema:
     """
     Apply validators to a schema.
     """
     for validator in validators:
         assert validator.sub_path is None, 'validator.sub_path is not yet supported'
         function = typing.cast(typing.Callable[..., Any], validator.function)
         if validator.mode == 'plain':
-            schema = {  # type: ignore[misc,assignment]
-                'type': 'function',
-                'mode': 'plain',
-                'function': function,
-            }
+            schema = core_schema.FunctionPlainSchema(
+                type='function',
+                mode='plain',
+                function=function,
+            )
         else:
-            schema = {
-                'type': 'function',
-                'mode': validator.mode,
-                'function': function,
-                'schema': schema,
-            }
+            schema = core_schema.FunctionSchema(
+                type='function',
+                mode=validator.mode,
+                function=function,
+                schema=schema,
+            )
     return schema
 
 
 class PydanticSchemaGenerationError(TypeError):
     pass
 
 
-def union_schema(union_type: Any) -> PydanticCoreSchema:
+def union_schema(union_type: Any) -> core_schema.Schema:
     """
     Generate schema for a Union.
     """
-    return {'type': 'union', 'choices': [generate_schema(arg) for arg in get_args(union_type)]}
+    args = get_args(union_type)
+    choices = []
+    nullable = False
+    for arg in args:
+        if arg is None or arg is NoneType:
+            nullable = True
+        else:
+            choices.append(generate_schema(arg))
+
+    if len(choices) == 1:
+        s = choices[0]
+    else:
+        s = core_schema.UnionSchema(type='union', choices=choices)
 
+    if nullable:
+        s = core_schema.NullableSchema(type='nullable', schema=s)
+    return s
 
-def literal_schema(literal_type: Any) -> PydanticCoreSchema:
+
+def literal_schema(literal_type: Any) -> core_schema.Schema:
     """
     Generate schema for a Literal.
     """
@@ -172,12 +189,12 @@ def literal_schema(literal_type: Any) -> PydanticCoreSchema:
     return {'type': 'literal', 'expected': list(expected)}
 
 
-def type_dict_schema(typed_dict: Any) -> TypedDictSchema:
+def type_dict_schema(typed_dict: Any) -> core_schema.TypedDictSchema:
     """
     Generate schema for a TypedDict.
     """
     required_keys: typing.Set[str] = getattr(typed_dict, '__required_keys__', set())
-    fields: typing.Dict[str, TypedDictField] = {}
+    fields: typing.Dict[str, core_schema.TypedDictField] = {}
 
     for field_name, field_type in typed_dict.__annotations__.items():
         required = field_name in required_keys
@@ -206,10 +223,10 @@ def type_dict_schema(typed_dict: Any) -> TypedDictSchema:
 
         fields[field_name] = {'schema': schema, 'required': required}
 
-    return {'type': 'typed-dict', 'fields': fields, 'extra_behavior': 'forbid'}
+    return core_schema.TypedDictSchema(type='typed-dict', fields=fields, extra_behavior='forbid')
 
 
-def generic_collection_schema(obj: Any) -> PydanticCoreSchema:
+def generic_collection_schema(obj: Any) -> core_schema.Schema:
     """
     Generate schema for List, Set etc. - where the schema includes `items_schema`
 
@@ -221,14 +238,14 @@ def generic_collection_schema(obj: Any) -> PydanticCoreSchema:
     }  # type: ignore[misc,return-value]
 
 
-def tuple_schema(tuple_type: Any) -> typing.Union[TupleVariableSchema, TuplePositionalSchema]:
+def tuple_schema(tuple_type: Any) -> core_schema.Schema:
     """
     Generate schema for a Tuple, e.g. `tuple[int, str]` or `tuple[int, ...]`.
     """
     params = get_args(tuple_type)
     if params[-1] is Ellipsis:
         if len(params) == 2:
-            sv: TupleVariableSchema = {'type': 'tuple', 'mode': 'variable', 'items_schema': generate_schema(params[0])}
+            sv = core_schema.TupleVariableSchema(type='tuple', mode='variable', items_schema=generate_schema(params[0]))
             return sv
 
         # not sure this case is valid in python, but may as well support it here since pydantic-core does
@@ -240,32 +257,32 @@ def tuple_schema(tuple_type: Any) -> typing.Union[TupleVariableSchema, TuplePosi
             'extra_schema': generate_schema(extra_schema),
         }
     else:
-        sp: TuplePositionalSchema = {
-            'type': 'tuple',
-            'mode': 'positional',
-            'items_schema': [generate_schema(p) for p in params],
-        }
+        sp = core_schema.TuplePositionalSchema(
+            type='tuple',
+            mode='positional',
+            items_schema=[generate_schema(p) for p in params],
+        )
         return sp
 
 
-def dict_schema(dict_type: Any) -> DictSchema:
+def dict_schema(dict_type: Any) -> core_schema.DictSchema:
     """
     Generate schema for a Dict, e.g. `dict[str, int]`.
     """
     arg0, arg1 = get_args(dict_type)
-    return {
-        'type': 'dict',
-        'keys_schema': generate_schema(arg0),
-        'values_schema': generate_schema(arg1),
-    }
+    return core_schema.DictSchema(
+        type='dict',
+        keys_schema=generate_schema(arg0),
+        values_schema=generate_schema(arg1),
+    )
 
 
-def type_schema(type_: Any) -> IsInstanceSchema:
+def type_schema(type_: Any) -> core_schema.IsInstanceSchema:
     """
     Generate schema for a Type, e.g. `Type[int]`.
     """
     type_param = get_args(type_)[0]
     if type_param == Any:
-        return {'type': 'is-instance', 'class_': type}
+        return core_schema.IsInstanceSchema(type='is-instance', class_=type)
     else:
-        return {'type': 'is-instance', 'class_': type_param}
+        return core_schema.IsInstanceSchema(type='is-instance', class_=type_param)
diff --git a/pydantic/_internal/model_construction.py b/pydantic/_internal/model_construction.py
@@ -3,9 +3,10 @@
 import sys
 import warnings
 from types import FunctionType
-from typing import TYPE_CHECKING, Any, Callable, Type, get_type_hints
+from typing import TYPE_CHECKING, Any, Callable, Type, Union, get_type_hints
 
 from pydantic_core import SchemaValidator
+from pydantic_core.schema_types import NewClassSchema, RecursiveReferenceSchema, Schema as PydanticCoreSchema
 
 from ..fields import FieldInfo, ModelPrivateAttr, PrivateAttr
 from ..utils import ClassAttribute, is_valid_identifier
@@ -69,10 +70,26 @@ class SelfType:
     This is used to identify a reference to the current model class, e.g. in recursive classes
     """
 
-    __slots__ = ('name',)
+    __slots__ = ('__pydantic_validation_schema__',)
 
-    def __init__(self, module_name: str, cls_name: str):
-        self.name = f'{module_name}.{cls_name}'
+    def __init__(self, schema: PydanticCoreSchema):
+        self.__pydantic_validation_schema__ = schema
+
+    def __call__(self) -> None:
+        """
+        This is here just to mollify typing._type_check which expects typing types
+        but will also accept callables
+        """
+        raise TypeError('SelfType cannot be called')
+
+    def __or__(self, right: Any) -> Any:
+        return Union[self, right]
+
+    def __ror__(self, left: Any) -> Any:
+        return Union[left, self]
+
+    def __repr__(self) -> str:
+        return f'SelfType[{self.__pydantic_validation_schema__}]'
 
 
 def complete_model_class(
@@ -98,7 +115,14 @@ def complete_model_class(
             base_globals = module.__dict__
 
     fields: dict[str, FieldInfo] = {}
-    for ann_name, ann_type in get_type_hints(cls, base_globals, {name: SelfType(module_name, name)}).items():
+    core_config = generate_config(cls.__config__)
+    model_ref = f'{module_name}.{name}'
+    self_schema = NewClassSchema(
+        type='new-class',
+        class_type=cls,
+        schema=RecursiveReferenceSchema(type='recursive-ref', schema_ref=model_ref),
+    )
+    for ann_name, ann_type in get_type_hints(cls, base_globals, {name: SelfType(self_schema)}).items():
         # TODO NotField
         if ann_name.startswith('_') or is_classvar(ann_type):
             continue
@@ -121,18 +145,17 @@ def complete_model_class(
             # 2. To avoid false positives in the NameError check above
             delattr(cls, ann_name)
 
-    core_config = generate_config(cls.__config__)
-    inner_schema = model_fields_schema(fields, validator_functions)
+    inner_schema = model_fields_schema(model_ref, fields, validator_functions)
     validator_functions.check_for_unused()
 
     cls.__fields__ = fields
     cls.__pydantic_validator__ = SchemaValidator(inner_schema, core_config)
-    cls.__pydantic_validation_schema__ = {
-        'type': 'new-class',
-        'class_type': cls,
-        'schema': inner_schema,
-        'config': core_config,
-    }
+    cls.__pydantic_validation_schema__ = NewClassSchema(
+        type='new-class',
+        class_type=cls,
+        schema=inner_schema,
+        config=core_config,
+    )
 
     # set __signature__ attr only for model class, but not for its instances
     cls.__signature__ = ClassAttribute('__signature__', generate_model_signature(cls.__init__, fields, cls.__config__))

diff --git a/pydantic/main.py b/pydantic/main.py
@@ -523,7 +523,7 @@ def update_forward_refs(cls, **localns: Any) -> None:
         """
         Try to update ForwardRefs on fields based on this Model and localns.
         """
-        raise NotImplementedError('TODO, also rename to model_rebuild')
+        raise RuntimeError('TODO, also rename to model_rebuild')
 
     def __iter__(self) -> 'TupleGenerator':
         """

diff --git a/pyproject.toml b/pyproject.toml
@@ -40,7 +40,7 @@ classifiers = [
 requires-python = '>=3.7'
 dependencies = [
     'typing-extensions>=4.1.0',
-    'pydantic-core>=0.1.0',
+    'pydantic-core>=0.2.1',
 ]
 optional-dependencies = { email = ['email-validator>=1.0.3'] }
 dynamic = ['version']

diff --git a/requirements/pyproject-all.txt b/requirements/pyproject-all.txt
@@ -10,5 +10,9 @@ email-validator==1.2.1
     # via pydantic (pyproject.toml)
 idna==3.3
     # via email-validator
-typing-extensions==4.3.0
+pydantic-core==0.2.1
     # via pydantic (pyproject.toml)
+typing-extensions==4.3.0
+    # via
+    #   pydantic (pyproject.toml)
+    #   pydantic-core
diff --git a/requirements/pyproject-min.txt b/requirements/pyproject-min.txt
@@ -4,5 +4,9 @@
 #
 #    pip-compile --output-file=requirements/pyproject-min.txt pyproject.toml
 #
-typing-extensions==4.3.0
+pydantic-core==0.2.1
     # via pydantic (pyproject.toml)
+typing-extensions==4.3.0
+    # via
+    #   pydantic (pyproject.toml)
+    #   pydantic-core