Use union for multiple allowed literal values

pydantic · Jul 13, 2019 · a6533a9 · a6533a9
1 parent c8008d1
commit a6533a9
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 12 deletions.
diff --git a/pydantic/fields.py b/pydantic/fields.py
@@ -197,10 +197,10 @@ def _populate_sub_fields(self) -> None:  # noqa: C901 (ignore complexity)
             # python 3.7 only, Pattern is a typing object but without sub fields
             return
         if is_literal_type(self.type_):
-            types_ = sorted(set(type(value) for value in literal_values(self.type_)), key=str)
-            if len(types_) > 1:
-                self.sub_fields = [self._create_sub_type(t, f'{self.name}_{display_as_type(t)}') for t in types_]
-            return
+            if len(literal_values(self.type_)) > 1:
+                self.type_ = Union[tuple(Literal[value] for value in literal_values(self.type_))]
+            else:
+                return
         origin = getattr(self.type_, '__origin__', None)
         if origin is None:
             # field is not "typing" object eg. Union, Dict, List etc.

diff --git a/pydantic/schema.py b/pydantic/schema.py
@@ -754,8 +754,10 @@ def field_singleton_schema(  # noqa: C901 (ignore complexity)
     if is_new_type(field_type):
         field_type = new_type_supertype(field_type)
     if is_literal_type(field_type):
-        # If there were distinct literal value types, field.sub_fields would not be falsy
-        field_type = type(literal_values(field_type)[0])
+        # If there were multiple literal values, field.sub_fields would not be falsy
+        literal_value = literal_values(field_type)[0]
+        field_type = type(literal_value)
+        f_schema['const'] = literal_value
     if issubclass(field_type, Enum):
         f_schema.update({'enum': [item.value for item in field_type]})
         # Don't return immediately, to allow adding specific types

diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -1450,9 +1450,9 @@ class Model(BaseModel):
 
     assert Model.schema() == {
         'properties': {
-            'a': {'title': 'A', 'type': 'integer'},
-            'b': {'title': 'B', 'type': 'string'},
-            'c': {'anyOf': [{'type': 'integer'}, {'type': 'string'}], 'title': 'C'},
+            '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'},
         },
         'required': ['a', 'b', 'c'],
         'title': 'Model',

diff --git a/tests/test_types.py b/tests/test_types.py
@@ -1698,8 +1698,14 @@ class Model(BaseModel):
     assert exc_info.value.errors() == [
         {
             'loc': ('a_or_b',),
-            'msg': "unexpected value; permitted: 'a', 'b'",
+            'msg': "unexpected value; permitted: 'a'",
             'type': 'value_error.const',
-            'ctx': {'given': 'c', 'permitted': ('a', 'b')},
-        }
+            'ctx': {'given': 'c', 'permitted': ('a',)},
+        },
+        {
+            'loc': ('a_or_b',),
+            'msg': "unexpected value; permitted: 'b'",
+            'type': 'value_error.const',
+            'ctx': {'given': 'c', 'permitted': ('b',)},
+        },
     ]