refactor(schema): support properly Literal in generated JSON schema (…

…#2348)

* test: improve example

* refactor: use enum for Literal in JSON schema

* test: update test with new schema generation

* chore: add change file

changes/1350-PrettyWood.md

@@ -0,0 +1 @@
+use `enum` for `typing.Literal` in JSON schema

pydantic/schema.py

@@ -1,5 +1,6 @@
 import re
 import warnings
+from collections import defaultdict
 from datetime import date, datetime, time, timedelta
 from decimal import Decimal
 from enum import Enum
@@ -64,12 +65,12 @@
 from .typing import (
     NONE_TYPES,
     ForwardRef,
+    all_literal_values,
     get_args,
     get_origin,
     is_callable_type,
     is_literal_type,
     is_namedtuple,
-    literal_values,
 )
 from .utils import ROOT_KEY, get_model, lenient_issubclass, sequence_like
 
@@ -789,19 +790,21 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
         f_schema['const'] = field.default
     field_type = field.type_
     if is_literal_type(field_type):
-        values = literal_values(field_type)
-        if len(values) > 1:
+        values = all_literal_values(field_type)
+
+        if len({v.__class__ for v in values}) > 1:
             return field_schema(
-                multivalue_literal_field_for_schema(values, field),
+                multitypes_literal_field_for_schema(values, field),
                 by_alias=by_alias,
                 model_name_map=model_name_map,
                 ref_prefix=ref_prefix,
                 ref_template=ref_template,
                 known_models=known_models,
             )
-        literal_value = values[0]
-        field_type = literal_value.__class__
-        f_schema['const'] = literal_value
+
+        # All values have the same type
+        field_type = values[0].__class__
+        f_schema['enum'] = list(values)
 
     if lenient_issubclass(field_type, Enum):
         enum_name = model_name_map[field_type]
@@ -859,10 +862,19 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
     raise ValueError(f'Value not declarable with JSON Schema, field: {field}')
 
 
-def multivalue_literal_field_for_schema(values: Tuple[Any, ...], field: ModelField) -> ModelField:
+def multitypes_literal_field_for_schema(values: Tuple[Any, ...], field: ModelField) -> ModelField:
+    """
+    To support `Literal` with values of different types, we split it into multiple `Literal` with same type
+    e.g. `Literal['qwe', 'asd', 1, 2]` becomes `Union[Literal['qwe', 'asd'], Literal[1, 2]]`
+    """
+    literal_distinct_types = defaultdict(list)
+    for v in values:
+        literal_distinct_types[v.__class__].append(v)
+    distinct_literals = (Literal[tuple(same_type_values)] for same_type_values in literal_distinct_types.values())
+
     return ModelField(
         name=field.name,
-        type_=Union[tuple(Literal[value] for value in values)],  # type: ignore
+        type_=Union[tuple(distinct_literals)],  # type: ignore
         class_validators=field.class_validators,
         model_config=field.model_config,
         default=field.default,

pydantic/typing.py

@@ -203,7 +203,7 @@ def get_args(tp: Type[Any]) -> Tuple[Any, ...]:
     'resolve_annotations',
     'is_callable_type',
     'is_literal_type',
-    'literal_values',
+    'all_literal_values',
     'is_namedtuple',
     'is_typeddict',
     'is_new_type',

tests/test_schema.py

@@ -1759,14 +1759,22 @@ class Model(BaseModel):
         a: Literal[1]
         b: Literal['a']
         c: Literal['a', 1]
+        d: Literal['a', Literal['b'], 1, 2]
 
     assert Model.schema() == {
         'properties': {
-            'a': {'title': 'A', 'type': 'integer', 'const': 1},
-            'b': {'title': 'B', 'type': 'string', 'const': 'a'},
-            'c': {'anyOf': [{'type': 'string', 'const': 'a'}, {'type': 'integer', 'const': 1}], 'title': 'C'},
+            'a': {'title': 'A', 'type': 'integer', 'enum': [1]},
+            'b': {'title': 'B', 'type': 'string', 'enum': ['a']},
+            'c': {'title': 'C', 'anyOf': [{'type': 'string', 'enum': ['a']}, {'type': 'integer', 'enum': [1]}]},
+            'd': {
+                'title': 'D',
+                'anyOf': [
+                    {'type': 'string', 'enum': ['a', 'b']},
+                    {'type': 'integer', 'enum': [1, 2]},
+                ],
+            },
         },
-        'required': ['a', 'b', 'c'],
+        'required': ['a', 'b', 'c', 'd'],
         'title': 'Model',
         'type': 'object',
     }