Skip to content

Commit

Permalink
Add enum and type to the JSON schema for single item literals (pydant…
Browse files Browse the repository at this point in the history 
…ic#8944)
  • Loading branch information
@dmontagu
dmontagu authored Mar 12, 2024
1 parent 3423213 commit 18d39fe
Show file tree
Hide file tree
Showing 8 changed files with 106 additions and 70 deletions.
18 changes: 8 additions & 10 deletions pydantic/json_schema.py
View file
Original file line number Diff line number Diff line change
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
View file
Original file line number Diff line number Diff line change
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
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1229,16 +1229,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 @@ -1311,11 +1323,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 @@ -1324,7 +1337,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 @@ -1790,20 +1803,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 @@ -1823,7 +1837,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 @@ -1841,7 +1855,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
View file
Original file line number Diff line number Diff line change
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
View file
Original file line number Diff line number Diff line change
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
View file
Original file line number Diff line number Diff line change
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
Loading

0 comments on commit 18d39fe

Please sign in to comment.