Generate schema for generic models (#2364)

pydantic/schema.py

@@ -374,7 +374,7 @@ def get_flat_models_from_field(field: ModelField, known_models: TypeModelSet) ->
     field_type = field.type_
     if lenient_issubclass(getattr(field_type, '__pydantic_model__', None), BaseModel):
         field_type = field_type.__pydantic_model__
-    if field.sub_fields:
+    if field.sub_fields and not lenient_issubclass(field_type, BaseModel):
         flat_models |= get_flat_models_from_fields(field.sub_fields, known_models=known_models)
     elif lenient_issubclass(field_type, BaseModel) and field_type not in known_models:
         flat_models |= get_flat_models_from_model(field_type, known_models=known_models)
@@ -769,7 +769,13 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
 
     definitions: Dict[str, Any] = {}
     nested_models: Set[str] = set()
-    if field.sub_fields:
+    field_type = field.type_
+
+    # Recurse into this field if it contains sub_fields and is NOT a
+    # BaseModel OR that BaseModel is a const
+    if field.sub_fields and (
+        (field.field_info and field.field_info.const) or not lenient_issubclass(field_type, BaseModel)
+    ):
         return field_singleton_sub_fields_schema(
             field.sub_fields,
             by_alias=by_alias,
@@ -779,16 +785,15 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
             ref_template=ref_template,
             known_models=known_models,
         )
-    if field.type_ is Any or field.type_.__class__ == TypeVar:
+    if field_type is Any or field_type.__class__ == TypeVar:
         return {}, definitions, nested_models  # no restrictions
-    if field.type_ in NONE_TYPES:
+    if field_type in NONE_TYPES:
         return {'type': 'null'}, definitions, nested_models
-    if is_callable_type(field.type_):
+    if is_callable_type(field_type):
         raise SkipField(f'Callable {field.name} was excluded from schema since JSON schema has no equivalent type.')
     f_schema: Dict[str, Any] = {}
     if field.field_info is not None and field.field_info.const:
         f_schema['const'] = field.default
-    field_type = field.type_
     if is_literal_type(field_type):
         values = all_literal_values(field_type)
 

tests/test_schema.py

@@ -30,6 +30,7 @@ from typing_extensions import Literal
 from pydantic import BaseModel, Extra, Field, ValidationError, conlist, conset, validator
 from pydantic.color import Color
 from pydantic.dataclasses import dataclass
+from pydantic.generics import GenericModel
 from pydantic.networks import AnyUrl, EmailStr, IPvAnyAddress, IPvAnyInterface, IPvAnyNetwork, NameEmail, stricturl
 from pydantic.schema import (
     get_flat_models_from_model,
@@ -82,6 +83,9 @@ except ImportError:
     email_validator = None
 
 
+T = TypeVar('T')
+
+
 def test_key():
     class ApplePie(BaseModel):
         """
@@ -2272,3 +2276,124 @@ def test_schema_for_generic_field():
         },
         'required': ['data', 'data1'],
     }
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 7), reason='schema generation for generic fields is not available in python < 3.7'
+)
+def test_nested_generic():
+    """
+    Test a nested BaseModel that is also a Generic
+    """
+
+    class Ref(BaseModel, Generic[T]):
+        uuid: str
+
+        def resolve(self) -> T:
+            ...
+
+    class Model(BaseModel):
+        ref: Ref['Model']  # noqa
+
+    assert Model.schema() == {
+        'title': 'Model',
+        'type': 'object',
+        'definitions': {
+            'Ref': {
+                'title': 'Ref',
+                'type': 'object',
+                'properties': {
+                    'uuid': {'title': 'Uuid', 'type': 'string'},
+                },
+                'required': ['uuid'],
+            },
+        },
+        'properties': {
+            'ref': {'$ref': '#/definitions/Ref'},
+        },
+        'required': ['ref'],
+    }
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 7), reason='schema generation for generic fields is not available in python < 3.7'
+)
+def test_nested_generic_model():
+    """
+    Test a nested GenericModel
+    """
+
+    class Box(GenericModel, Generic[T]):
+        uuid: str
+        data: T
+
+    class Model(BaseModel):
+        box_str: Box[str]
+        box_int: Box[int]
+
+    assert Model.schema() == {
+        'title': 'Model',
+        'type': 'object',
+        'definitions': {
+            'Box_str_': Box[str].schema(),
+            'Box_int_': Box[int].schema(),
+        },
+        'properties': {
+            'box_str': {'$ref': '#/definitions/Box_str_'},
+            'box_int': {'$ref': '#/definitions/Box_int_'},
+        },
+        'required': ['box_str', 'box_int'],
+    }
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 7), reason='schema generation for generic fields is not available in python < 3.7'
+)
+def test_complex_nested_generic():
+    """
+    Handle a union of a generic.
+    """
+
+    class Ref(BaseModel, Generic[T]):
+        uuid: str
+
+        def resolve(self) -> T:
+            ...
+
+    class Model(BaseModel):
+        uuid: str
+        model: Union[Ref['Model'], 'Model']  # noqa
+
+        def resolve(self) -> 'Model':  # noqa
+            ...
+
+    Model.update_forward_refs()
+
+    assert Model.schema() == {
+        'definitions': {
+            'Model': {
+                'title': 'Model',
+                'type': 'object',
+                'properties': {
+                    'uuid': {'title': 'Uuid', 'type': 'string'},
+                    'model': {
+                        'title': 'Model',
+                        'anyOf': [
+                            {'$ref': '#/definitions/Ref'},
+                            {'$ref': '#/definitions/Model'},
+                        ],
+                    },
+                },
+                'required': ['uuid', 'model'],
+            },
+            'Ref': {
+                'title': 'Ref',
+                'type': 'object',
+                'properties': {
+                    'uuid': {'title': 'Uuid', 'type': 'string'},
+                },
+                'required': ['uuid'],
+            },
+        },
+        '$ref': '#/definitions/Model',
+    }