Add type key to enum json schema (#6243)

pydantic · Jun 25, 2023 · ba8ee95 · ba8ee95
1 parent 44e28e7
commit ba8ee95
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 28 deletions.
diff --git a/pydantic/_internal/_dataclasses.py b/pydantic/_internal/_dataclasses.py
@@ -161,7 +161,7 @@ def is_builtin_dataclass(_cls: type[Any]) -> TypeGuard[type[StandardDataclass]]:
 
     we check that
     - `_cls` is a dataclass
-    - `_cls` is not a processed pydantic dataclass (with a basemodel attached)
+    - `_cls` is not a processed pydantic dataclass (with a `BaseModel` attached)
     - `_cls` is not a pydantic dataclass inheriting directly from a stdlib dataclass
     e.g.
     ```py

diff --git a/pydantic/json_schema.py b/pydantic/json_schema.py
@@ -556,7 +556,21 @@ def literal_schema(self, schema: core_schema.LiteralSchema) -> JsonSchemaValue:
         if len(expected) == 1:
             return {'const': expected[0]}
         else:
-            return {'enum': expected}
+            types = {type(e) for e in expected}
+            if types == {str}:
+                return {'enum': expected, 'type': 'string'}
+            elif types == {int}:
+                return {'enum': expected, 'type': 'integer'}
+            elif types == {float}:
+                return {'enum': expected, 'type': 'number'}
+            elif types == {bool}:
+                return {'enum': expected, 'type': 'boolean'}
+            elif types == {list}:
+                return {'enum': expected, 'type': 'array'}
+            # there is not None case because if it's mixed it hits the final `else`
+            # if it's a single Literal[None] then it becomes a `const` schema above
+            else:
+                return {'enum': expected}
 
     def is_instance_schema(self, schema: core_schema.IsInstanceSchema) -> JsonSchemaValue:
         """Returns a schema that checks if a value is an instance of a class, equivalent to Python's `isinstance`

diff --git a/tests/test_json_schema.py b/tests/test_json_schema.py
@@ -1,6 +1,7 @@
 import json
 import math
 import re
+import sys
 import typing
 from datetime import date, datetime, time, timedelta
 from decimal import Decimal
@@ -252,19 +253,19 @@ class Model(BaseModel):
 
     # insert_assert(Model.model_json_schema())
     assert Model.model_json_schema() == {
-        'type': 'object',
+        '$defs': {
+            'BarEnum': {'enum': [1, 2], 'title': 'BarEnum', 'type': 'integer'},
+            'FooEnum': {'enum': ['f', 'b'], 'title': 'FooEnum', 'type': 'string'},
+            'SpamEnum': {'enum': ['f', 'b'], 'title': 'SpamEnum', 'type': 'string'},
+        },
         'properties': {
             'foo': {'$ref': '#/$defs/FooEnum'},
             'bar': {'$ref': '#/$defs/BarEnum'},
             'spam': {'allOf': [{'$ref': '#/$defs/SpamEnum'}], 'default': None},
         },
         'required': ['foo', 'bar'],
         'title': 'Model',
-        '$defs': {
-            'BarEnum': {'enum': [1, 2], 'title': 'BarEnum', 'type': 'integer'},
-            'SpamEnum': {'enum': ['f', 'b'], 'title': 'SpamEnum', 'type': 'string'},
-            'FooEnum': {'enum': ['f', 'b'], 'title': 'FooEnum'},
-        },
+        'type': 'object',
     }
 
 
@@ -2066,11 +2067,66 @@ class Model(BaseModel):
         kind: Literal[MyEnum.FOO]
         other: Literal[MyEnum.FOO, MyEnum.BAR]
 
+    # insert_assert(Model.model_json_schema())
     assert Model.model_json_schema() == {
+        'properties': {
+            'kind': {'const': 'foo', 'title': 'Kind'},
+            'other': {'enum': ['foo', 'bar'], 'title': 'Other', 'type': 'string'},
+        },
+        'required': ['kind', 'other'],
+        'title': 'Model',
+        'type': 'object',
+    }
+
+
+@pytest.mark.skipif(sys.version_info[:2] == (3, 8), reason="ListEnum doesn't work in 3.8")
+def test_literal_types() -> None:
+    """Test that we properly add `type` to json schema enums when there is a single type."""
+
+    # for float and array we use an Enum because Literal can only accept str, int, bool or None
+    class FloatEnum(float, Enum):
+        a = 123.0
+        b = 123.1
+
+    class ListEnum(List[int], Enum):
+        a = [123]
+        b = [456]
+
+    class Model(BaseModel):
+        str_literal: Literal['foo', 'bar']
+        int_literal: Literal[123, 456]
+        float_literal: FloatEnum
+        bool_literal: Literal[True, False]
+        none_literal: Literal[None]  # ends up as a const since there's only 1
+        list_literal: ListEnum
+        mixed_literal: Literal[123, 'abc']
+
+    # insert_assert(Model.model_json_schema())
+    assert Model.model_json_schema() == {
+        '$defs': {
+            'FloatEnum': {'enum': [123.0, 123.1], 'title': 'FloatEnum', 'type': 'numeric'},
+            'ListEnum': {'enum': [[123], [456]], 'title': 'ListEnum', 'type': 'array'},
+        },
+        'properties': {
+            'str_literal': {'enum': ['foo', 'bar'], 'title': 'Str Literal', 'type': 'string'},
+            'int_literal': {'enum': [123, 456], 'title': 'Int Literal', 'type': 'integer'},
+            'float_literal': {'$ref': '#/$defs/FloatEnum'},
+            'bool_literal': {'enum': [True, False], 'title': 'Bool Literal', 'type': 'boolean'},
+            'none_literal': {'const': None, 'title': 'None Literal'},
+            'list_literal': {'$ref': '#/$defs/ListEnum'},
+            'mixed_literal': {'enum': [123, 'abc'], 'title': 'Mixed Literal'},
+        },
+        'required': [
+            'str_literal',
+            'int_literal',
+            'float_literal',
+            'bool_literal',
+            'none_literal',
+            'list_literal',
+            'mixed_literal',
+        ],
         'title': 'Model',
         'type': 'object',
-        'properties': {'kind': {'const': 'foo', 'title': 'Kind'}, 'other': {'enum': ['foo', 'bar'], 'title': 'Other'}},
-        'required': ['kind', 'other'],
     }
 
 
@@ -2358,18 +2414,20 @@ class ExampleEnum(Enum):
     class Example(BaseModel):
         example: ExampleEnum
 
+    # insert_assert(Example.model_json_schema())
     assert Example.model_json_schema() == {
-        'title': 'Example',
-        'type': 'object',
-        'properties': {'example': {'$ref': '#/$defs/ExampleEnum'}},
-        'required': ['example'],
         '$defs': {
             'ExampleEnum': {
-                'title': 'ExampleEnum',
                 'description': 'This is a test description.',
                 'enum': ['GT', 'LT', 'GE', 'LE', 'ML', 'MO', 'RE'],
+                'title': 'ExampleEnum',
+                'type': 'string',
             }
         },
+        'properties': {'example': {'$ref': '#/$defs/ExampleEnum'}},
+        'required': ['example'],
+        'title': 'Example',
+        'type': 'object',
     }
 
 
@@ -2839,7 +2897,7 @@ class Model(BaseModel):
 
     # insert_assert(Model.model_json_schema())
     assert Model.model_json_schema() == {
-        '$defs': {'CustomType': {'enum': ['a', 'b'], 'title': 'CustomType'}},
+        '$defs': {'CustomType': {'enum': ['a', 'b'], 'title': 'CustomType', 'type': 'string'}},
         'properties': {
             'data0': {
                 'anyOf': [{'type': 'string'}, {'items': {'type': 'string'}, 'type': 'array'}],
@@ -3055,6 +3113,7 @@ class Lizard(BaseModel):
     class Model(BaseModel):
         pet: Union[Cat, Dog, Lizard] = Field(..., discriminator='pet_type')
 
+    # insert_assert(Model.model_json_schema())
     assert Model.model_json_schema() == {
         '$defs': {
             'Cat': {
@@ -3070,7 +3129,7 @@ class Model(BaseModel):
                 'type': 'object',
             },
             'Lizard': {
-                'properties': {'pet_type': {'enum': ['reptile', 'lizard'], 'title': 'Pet Type'}},
+                'properties': {'pet_type': {'enum': ['reptile', 'lizard'], 'title': 'Pet Type', 'type': 'string'}},
                 'required': ['pet_type'],
                 'title': 'Lizard',
                 'type': 'object',
@@ -3087,11 +3146,7 @@ class Model(BaseModel):
                     },
                     'propertyName': 'pet_type',
                 },
-                'oneOf': [
-                    {'$ref': '#/$defs/Cat'},
-                    {'$ref': '#/$defs/Dog'},
-                    {'$ref': '#/$defs/Lizard'},
-                ],
+                'oneOf': [{'$ref': '#/$defs/Cat'}, {'$ref': '#/$defs/Dog'}, {'$ref': '#/$defs/Lizard'}],
                 'title': 'Pet',
             }
         },
@@ -3114,6 +3169,7 @@ class Lizard(BaseModel):
     class Model(BaseModel):
         pet: Annotated[Union[Cat, Dog, Lizard], Field(..., discriminator='pet_type')]
 
+    # insert_assert(Model.model_json_schema())
     assert Model.model_json_schema() == {
         '$defs': {
             'Cat': {
@@ -3129,7 +3185,7 @@ class Model(BaseModel):
                 'type': 'object',
             },
             'Lizard': {
-                'properties': {'pet_type': {'enum': ['reptile', 'lizard'], 'title': 'Pet Type'}},
+                'properties': {'pet_type': {'enum': ['reptile', 'lizard'], 'title': 'Pet Type', 'type': 'string'}},
                 'required': ['pet_type'],
                 'title': 'Lizard',
                 'type': 'object',
@@ -3146,11 +3202,7 @@ class Model(BaseModel):
                     },
                     'propertyName': 'pet_type',
                 },
-                'oneOf': [
-                    {'$ref': '#/$defs/Cat'},
-                    {'$ref': '#/$defs/Dog'},
-                    {'$ref': '#/$defs/Lizard'},
-                ],
+                'oneOf': [{'$ref': '#/$defs/Cat'}, {'$ref': '#/$defs/Dog'}, {'$ref': '#/$defs/Lizard'}],
                 'title': 'Pet',
             }
         },