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

@@ -0,0 +1 @@
+Support custom title, description and default in schema of enums

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}')
 

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'