From 3aacec4e1746d2446a50bc8660f659e46a296ea2 Mon Sep 17 00:00:00 2001 From: PrettyWood Date: Sun, 18 Oct 2020 21:23:06 +0200 Subject: [PATCH] feat(schema): support custom title, description and default for enums (#1749) * refactor(schema): put schema data from fieldinfo in dedicated function * feat(schema): support custom title, description and default for enums closes #1748 * refactor: replace $ref by allOf + $ref to be supported by doc generation tools * fix: do not set title by default for enums * refactor: make code more explicit * fix: run linter --- changes/1748-PrettyWood.md | 1 + pydantic/schema.py | 57 +++++++++++++++------- tests/test_schema.py | 97 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 137 insertions(+), 18 deletions(-) create mode 100644 changes/1748-PrettyWood.md diff --git a/changes/1748-PrettyWood.md b/changes/1748-PrettyWood.md new file mode 100644 index 0000000..e2d366e --- /dev/null +++ b/changes/1748-PrettyWood.md @@ -0,0 +1 @@ +Support custom title, description and default in schema of enums \ No newline at end of file diff --git a/pydantic/schema.py b/pydantic/schema.py index 15b5066..a7c105f 100644 --- a/pydantic/schema.py +++ b/pydantic/schema.py @@ -149,6 +149,30 @@ def model_schema( return m_schema +def get_field_info_schema(field: ModelField) -> Tuple[Dict[str, Any], bool]: + schema_overrides = False + + # If no title is explicitly set, we don't set title in the schema for enums. + # The behaviour is the same as `BaseModel` reference, where the default title + # is in the definitions part of the schema. + schema: Dict[str, Any] = {} + if field.field_info.title or not lenient_issubclass(field.type_, Enum): + schema['title'] = field.field_info.title or field.alias.title().replace('_', ' ') + + if field.field_info.title: + schema_overrides = True + + if field.field_info.description: + schema['description'] = field.field_info.description + schema_overrides = True + + if not field.required and not field.field_info.const and field.default is not None: + schema['default'] = encode_default(field.default) + schema_overrides = True + + return schema, schema_overrides + + def field_schema( field: ModelField, *, @@ -172,18 +196,7 @@ def field_schema( :return: tuple of the schema for this field and additional definitions """ ref_prefix = ref_prefix or default_prefix - schema_overrides = False - s = dict(title=field.field_info.title or field.alias.title().replace('_', ' ')) - if field.field_info.title: - schema_overrides = True - - if field.field_info.description: - s['description'] = field.field_info.description - schema_overrides = True - - if not field.required and not field.field_info.const and field.default is not None: - s['default'] = encode_default(field.default) - schema_overrides = True + s, schema_overrides = get_field_info_schema(field) validation_schema = get_field_schema_validations(field) if validation_schema: @@ -228,6 +241,11 @@ def get_field_schema_validations(field: ModelField) -> Dict[str, Any]: a Pydantic ``FieldInfo`` with validation arguments. """ f_schema: Dict[str, Any] = {} + + if lenient_issubclass(field.type_, Enum): + # schema is already updated by `enum_process_schema` + return f_schema + if lenient_issubclass(field.type_, (str, bytes)): for attr_name, t, keyword in _str_types_attrs: attr = getattr(field.field_info, attr_name, None) @@ -651,6 +669,11 @@ def add_field_type_to_schema(field_type: Any, schema: Dict[str, Any]) -> None: break +def get_schema_ref(ref_name: str, schema_overrides: bool) -> Dict[str, Any]: + schema_ref = {'$ref': ref_name} + return {'allOf': [schema_ref]} if schema_overrides else schema_ref + + def field_singleton_schema( # noqa: C901 (ignore complexity) field: ModelField, *, @@ -703,7 +726,8 @@ def field_singleton_schema( # noqa: C901 (ignore complexity) if lenient_issubclass(field_type, Enum): enum_name = normalize_name(field_type.__name__) - f_schema = {'$ref': ref_prefix + enum_name} + f_schema, schema_overrides = get_field_info_schema(field) + f_schema.update(get_schema_ref(ref_prefix + enum_name, schema_overrides)) definitions[enum_name] = enum_process_schema(field_type) else: add_field_type_to_schema(field_type, f_schema) @@ -734,11 +758,8 @@ def field_singleton_schema( # noqa: C901 (ignore complexity) nested_models.update(sub_nested_models) else: nested_models.add(model_name) - schema_ref = {'$ref': ref_prefix + model_name} - if not schema_overrides: - return schema_ref, definitions, nested_models - else: - return {'allOf': [schema_ref]}, definitions, nested_models + schema_ref = get_schema_ref(ref_prefix + model_name, schema_overrides) + return schema_ref, definitions, nested_models raise ValueError(f'Value not declarable with JSON Schema, field: {field}') diff --git a/tests/test_schema.py b/tests/test_schema.py index b08690e..ef5d2b0 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -250,6 +250,103 @@ def test_enum_modify_schema(): } +def test_enum_schema_custom_field(): + class FooBarEnum(str, Enum): + foo = 'foo' + bar = 'bar' + + class Model(BaseModel): + pika: FooBarEnum = Field(alias='pikalias', title='Pikapika!', description='Pika is definitely the best!') + bulbi: FooBarEnum = Field('foo', alias='bulbialias', title='Bulbibulbi!', description='Bulbi is not...') + cara: FooBarEnum + + assert Model.schema() == { + 'definitions': { + 'FooBarEnum': { + 'description': 'An enumeration.', + 'enum': ['foo', 'bar'], + 'title': 'FooBarEnum', + 'type': 'string', + } + }, + 'properties': { + 'pikalias': { + 'allOf': [{'$ref': '#/definitions/FooBarEnum'}], + 'description': 'Pika is definitely the best!', + 'title': 'Pikapika!', + }, + 'bulbialias': { + 'allOf': [{'$ref': '#/definitions/FooBarEnum'}], + 'description': 'Bulbi is not...', + 'title': 'Bulbibulbi!', + 'default': 'foo', + }, + 'cara': {'$ref': '#/definitions/FooBarEnum'}, + }, + 'required': ['pikalias', 'cara'], + 'title': 'Model', + 'type': 'object', + } + + +def test_enum_and_model_have_same_behaviour(): + class Names(str, Enum): + rick = 'Rick' + morty = 'Morty' + summer = 'Summer' + + class Pika(BaseModel): + a: str + + class Foo(BaseModel): + enum: Names + titled_enum: Names = Field( + ..., + title='Title of enum', + description='Description of enum', + ) + model: Pika + titled_model: Pika = Field( + ..., + title='Title of model', + description='Description of model', + ) + + assert Foo.schema() == { + 'definitions': { + 'Pika': { + 'properties': {'a': {'title': 'A', 'type': 'string'}}, + 'required': ['a'], + 'title': 'Pika', + 'type': 'object', + }, + 'Names': { + 'description': 'An enumeration.', + 'enum': ['Rick', 'Morty', 'Summer'], + 'title': 'Names', + 'type': 'string', + }, + }, + 'properties': { + 'enum': {'$ref': '#/definitions/Names'}, + 'model': {'$ref': '#/definitions/Pika'}, + 'titled_enum': { + 'allOf': [{'$ref': '#/definitions/Names'}], + 'description': 'Description of enum', + 'title': 'Title of enum', + }, + 'titled_model': { + 'allOf': [{'$ref': '#/definitions/Pika'}], + 'description': 'Description of model', + 'title': 'Title of model', + }, + }, + 'required': ['enum', 'titled_enum', 'model', 'titled_model'], + 'title': 'Foo', + 'type': 'object', + } + + def test_json_schema(): class Model(BaseModel): a = b'foobar'