Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add enum and type to the JSON schema for single item literals #8944

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
18 changes: 8 additions & 10 deletions pydantic/json_schema.py
Expand Up @@ -732,24 +732,22 @@ def literal_schema(self, schema: core_schema.LiteralSchema) -> JsonSchemaValue:
# jsonify the expected values
expected = [to_jsonable_python(v) for v in expected]

result: dict[str, Any] = {'enum': expected}
if len(expected) == 1:
return {'const': expected[0]}
result['const'] = expected[0]

types = {type(e) for e in expected}
if types == {str}:
return {'enum': expected, 'type': 'string'}
result['type'] = 'string'
elif types == {int}:
return {'enum': expected, 'type': 'integer'}
result['type'] = 'integer'
elif types == {float}:
return {'enum': expected, 'type': 'number'}
result['type'] = 'number'
elif types == {bool}:
return {'enum': expected, 'type': 'boolean'}
result['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}
result['type'] = 'array'
return result

def is_instance_schema(self, schema: core_schema.IsInstanceSchema) -> JsonSchemaValue:
"""Handles JSON schema generation for a core schema that checks if a value is an instance of a class.
Expand Down
15 changes: 13 additions & 2 deletions tests/test_dataclasses.py
Expand Up @@ -1353,6 +1353,7 @@ class Top:

t = Top(sub=A(l='a'))
assert isinstance(t, Top)
# insert_assert(model_json_schema(Top))
assert model_json_schema(Top) == {
'title': 'Top',
'type': 'object',
Expand All @@ -1365,8 +1366,18 @@ class Top:
},
'required': ['sub'],
'$defs': {
'A': {'properties': {'l': {'const': 'a', 'title': 'L'}}, 'required': ['l'], 'title': 'A', 'type': 'object'},
'B': {'properties': {'l': {'const': 'b', 'title': 'L'}}, 'required': ['l'], 'title': 'B', 'type': 'object'},
'A': {
'properties': {'l': {'const': 'a', 'enum': ['a'], 'title': 'L', 'type': 'string'}},
'required': ['l'],
'title': 'A',
'type': 'object',
},
'B': {
'properties': {'l': {'const': 'b', 'enum': ['b'], 'title': 'L', 'type': 'string'}},
'required': ['l'],
'title': 'B',
'type': 'object',
},
},
}

Expand Down
34 changes: 24 additions & 10 deletions tests/test_discriminated_union.py
Expand Up @@ -1235,16 +1235,28 @@ class TestModel(BaseModel):
},
'UnionModel1': {
'properties': {
'type': {'const': 1, 'default': 1, 'title': 'Type'},
'other': {'const': 'UnionModel1', 'default': 'UnionModel1', 'title': 'Other'},
'type': {'const': 1, 'default': 1, 'enum': [1], 'title': 'Type', 'type': 'integer'},
'other': {
'const': 'UnionModel1',
'default': 'UnionModel1',
'enum': ['UnionModel1'],
'title': 'Other',
'type': 'string',
},
},
'title': 'UnionModel1',
'type': 'object',
},
'UnionModel2': {
'properties': {
'type': {'const': 2, 'default': 2, 'title': 'Type'},
'other': {'const': 'UnionModel2', 'default': 'UnionModel2', 'title': 'Other'},
'type': {'const': 2, 'default': 2, 'enum': [2], 'title': 'Type', 'type': 'integer'},
'other': {
'const': 'UnionModel2',
'default': 'UnionModel2',
'enum': ['UnionModel2'],
'title': 'Other',
'type': 'string',
},
},
'title': 'UnionModel2',
'type': 'object',
Expand Down Expand Up @@ -1317,11 +1329,12 @@ class Model(BaseModel):
pet: Sequence[Pet]
n: int

# insert_assert(Model.model_json_schema())
assert Model.model_json_schema() == {
'$defs': {
'Cat': {
'properties': {
'pet_type': {'const': 'cat', 'title': 'Pet Type'},
'pet_type': {'const': 'cat', 'enum': ['cat'], 'title': 'Pet Type', 'type': 'string'},
'meows': {'title': 'Meows', 'type': 'integer'},
},
'required': ['pet_type', 'meows'],
Expand All @@ -1330,7 +1343,7 @@ class Model(BaseModel):
},
'Dog': {
'properties': {
'pet_type': {'const': 'dog', 'title': 'Pet Type'},
'pet_type': {'const': 'dog', 'enum': ['dog'], 'title': 'Pet Type', 'type': 'string'},
'barks': {'title': 'Barks', 'type': 'number'},
},
'required': ['pet_type', 'barks'],
Expand Down Expand Up @@ -1796,20 +1809,21 @@ class SubModel(MyModel):
blending: float

MyModel.model_rebuild()
# insert_assert(MyModel.model_json_schema())
assert MyModel.model_json_schema() == {
'$defs': {
'Step_A': {
'properties': {
'count': {'title': 'Count', 'type': 'integer'},
'type': {'const': 'stepA', 'title': 'Type'},
'type': {'const': 'stepA', 'enum': ['stepA'], 'title': 'Type', 'type': 'string'},
},
'required': ['type', 'count'],
'title': 'Step_A',
'type': 'object',
},
'Step_B': {
'properties': {
'type': {'const': 'stepB', 'title': 'Type'},
'type': {'const': 'stepB', 'enum': ['stepB'], 'title': 'Type', 'type': 'string'},
'value': {'title': 'Value', 'type': 'number'},
},
'required': ['type', 'value'],
Expand All @@ -1829,7 +1843,7 @@ class SubModel(MyModel):
'title': 'Steps',
},
'sub_models': {'items': {'$ref': '#/$defs/SubModel'}, 'title': 'Sub Models', 'type': 'array'},
'type': {'const': 'mixed', 'title': 'Type'},
'type': {'const': 'mixed', 'enum': ['mixed'], 'title': 'Type', 'type': 'string'},
},
'required': ['type', 'sub_models', 'blending'],
'title': 'SubModel',
Expand All @@ -1847,7 +1861,7 @@ class SubModel(MyModel):
'title': 'Steps',
},
'sub_models': {'items': {'$ref': '#/$defs/SubModel'}, 'title': 'Sub Models', 'type': 'array'},
'type': {'const': 'mixed', 'title': 'Type'},
'type': {'const': 'mixed', 'enum': ['mixed'], 'title': 'Type', 'type': 'string'},
},
'required': ['type', 'sub_models'],
'title': 'MyModel',
Expand Down
5 changes: 3 additions & 2 deletions tests/test_edge_cases.py
Expand Up @@ -2665,11 +2665,12 @@ class Outer(BaseModel):
validated = Outer.model_validate({'a': {'kind': '1', 'two': None}, 'b': {'kind': '2', 'one': None}})
assert validated == Outer(a=Root1(root=Model1(two=None)), b=Root2(root=Model2(one=None)))

# insert_assert(Outer.model_json_schema())
assert Outer.model_json_schema() == {
'$defs': {
'Model1': {
'properties': {
'kind': {'const': '1', 'default': '1', 'title': 'Kind'},
'kind': {'const': '1', 'default': '1', 'enum': ['1'], 'title': 'Kind', 'type': 'string'},
'two': {'anyOf': [{'$ref': '#/$defs/Model2'}, {'type': 'null'}]},
},
'required': ['two'],
Expand All @@ -2678,7 +2679,7 @@ class Outer(BaseModel):
},
'Model2': {
'properties': {
'kind': {'const': '2', 'default': '2', 'title': 'Kind'},
'kind': {'const': '2', 'default': '2', 'enum': ['2'], 'title': 'Kind', 'type': 'string'},
'one': {'anyOf': [{'$ref': '#/$defs/Model1'}, {'type': 'null'}]},
},
'required': ['one'],
Expand Down
5 changes: 3 additions & 2 deletions tests/test_forward_ref.py
Expand Up @@ -540,6 +540,7 @@ class Dog(BaseModel):
# Ensure the rebuild has happened automatically despite validation failure
assert module.Pet.__pydantic_complete__ is True

# insert_assert(module.Pet.model_json_schema())
assert module.Pet.model_json_schema() == {
'title': 'Pet',
'required': ['pet'],
Expand All @@ -555,13 +556,13 @@ class Dog(BaseModel):
'Cat': {
'title': 'Cat',
'type': 'object',
'properties': {'type': {'const': 'cat', 'title': 'Type'}},
'properties': {'type': {'const': 'cat', 'enum': ['cat'], 'title': 'Type', 'type': 'string'}},
'required': ['type'],
},
'Dog': {
'title': 'Dog',
'type': 'object',
'properties': {'type': {'const': 'dog', 'title': 'Type'}},
'properties': {'type': {'const': 'dog', 'enum': ['dog'], 'title': 'Type', 'type': 'string'}},
'required': ['type'],
},
},
Expand Down
6 changes: 4 additions & 2 deletions tests/test_generics.py
Expand Up @@ -2433,8 +2433,9 @@ class Model(BaseModel, Generic[T]):
}
]

# insert_assert(Model[MyEnum].model_json_schema())
assert Model[MyEnum].model_json_schema() == {
'$defs': {'MyEnum': {'const': 1, 'title': 'MyEnum'}},
'$defs': {'MyEnum': {'const': 1, 'enum': [1], 'title': 'MyEnum', 'type': 'integer'}},
'properties': {'x': {'$ref': '#/$defs/MyEnum'}},
'required': ['x'],
'title': 'Model[test_generic_enum_bound.<locals>.MyEnum]',
Expand Down Expand Up @@ -2484,8 +2485,9 @@ class Model(BaseModel, Generic[T]):
}
]

# insert_assert(Model[MyEnum].model_json_schema())
assert Model[MyEnum].model_json_schema() == {
'$defs': {'MyEnum': {'const': 1, 'title': 'MyEnum', 'type': 'integer'}},
'$defs': {'MyEnum': {'const': 1, 'enum': [1], 'title': 'MyEnum', 'type': 'integer'}},
'properties': {'x': {'$ref': '#/$defs/MyEnum'}},
'required': ['x'],
'title': 'Model[test_generic_intenum_bound.<locals>.MyEnum]',
Expand Down