From 82fb6ebfc9bebcd8793abed77a86e1a3929d868f Mon Sep 17 00:00:00 2001 From: Danielle Madeley Date: Sun, 9 May 2021 18:23:56 +1000 Subject: [PATCH] Generate schema for generic models (#2364) --- pydantic/schema.py | 17 +++--- tests/test_schema.py | 125 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 6 deletions(-) diff --git a/pydantic/schema.py b/pydantic/schema.py index 4dc47f0..1fa1b8b 100644 --- a/pydantic/schema.py +++ b/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) diff --git a/tests/test_schema.py b/tests/test_schema.py index 29a7e56..b386cb6 100644 --- a/tests/test_schema.py +++ b/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', + }