Skip to content

Commit

Permalink
Fix regression in core schema generation for indirect definition refe…
Browse files Browse the repository at this point in the history
…rences (#8702)

Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
  • Loading branch information
dmontagu and sydney-runkle committed Feb 2, 2024
1 parent 8af855d commit 8898b8f
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 1 deletion.
2 changes: 1 addition & 1 deletion pydantic/_internal/_core_utils.py
Expand Up @@ -228,7 +228,7 @@ def _handle_ser_schemas(self, ser_schema: core_schema.SerSchema, f: Walk) -> cor
def handle_definitions_schema(self, schema: core_schema.DefinitionsSchema, f: Walk) -> core_schema.CoreSchema:
new_definitions: list[core_schema.CoreSchema] = []
for definition in schema['definitions']:
if 'schema_ref' and 'ref' in definition:
if 'schema_ref' in definition and 'ref' in definition:
# This indicates a purposely indirect reference
# We want to keep such references around for implications related to JSON schema, etc.:
new_definitions.append(definition)
Expand Down
88 changes: 88 additions & 0 deletions tests/test_discriminated_union.py
Expand Up @@ -1763,3 +1763,91 @@ class CreateObjectDto(BaseModel):
'type': 'object',
},
}


def test_nested_discriminator() -> None:
"""
The exact details of the JSON schema produced are not necessarily important; the test was added in response to a
regression that caused the inner union to lose its discriminator. Even if the schema changes, the important
thing is that the core schema (and therefore JSON schema) produced has an actual discriminated union in it.
For more context, see: https://github.com/pydantic/pydantic/issues/8688.
"""

class Step_A(BaseModel):
type: Literal['stepA']
count: int

class Step_B(BaseModel):
type: Literal['stepB']
value: float

class MyModel(BaseModel):
type: Literal['mixed']
sub_models: List['SubModel']
steps: Union[Step_A, Step_B] = Field(
default=None,
discriminator='type',
)

class SubModel(MyModel):
type: Literal['mixed']
blending: float

MyModel.model_rebuild()
assert MyModel.model_json_schema() == {
'$defs': {
'Step_A': {
'properties': {
'count': {'title': 'Count', 'type': 'integer'},
'type': {'const': 'stepA', 'title': 'Type'},
},
'required': ['type', 'count'],
'title': 'Step_A',
'type': 'object',
},
'Step_B': {
'properties': {
'type': {'const': 'stepB', 'title': 'Type'},
'value': {'title': 'Value', 'type': 'number'},
},
'required': ['type', 'value'],
'title': 'Step_B',
'type': 'object',
},
'SubModel': {
'properties': {
'blending': {'title': 'Blending', 'type': 'number'},
'steps': {
'default': None,
'discriminator': {
'mapping': {'stepA': '#/$defs/Step_A', 'stepB': '#/$defs/Step_B'},
'propertyName': 'type',
},
'oneOf': [{'$ref': '#/$defs/Step_A'}, {'$ref': '#/$defs/Step_B'}],
'title': 'Steps',
},
'sub_models': {'items': {'$ref': '#/$defs/SubModel'}, 'title': 'Sub Models', 'type': 'array'},
'type': {'const': 'mixed', 'title': 'Type'},
},
'required': ['type', 'sub_models', 'blending'],
'title': 'SubModel',
'type': 'object',
},
},
'properties': {
'steps': {
'default': None,
'discriminator': {
'mapping': {'stepA': '#/$defs/Step_A', 'stepB': '#/$defs/Step_B'},
'propertyName': 'type',
},
'oneOf': [{'$ref': '#/$defs/Step_A'}, {'$ref': '#/$defs/Step_B'}],
'title': 'Steps',
},
'sub_models': {'items': {'$ref': '#/$defs/SubModel'}, 'title': 'Sub Models', 'type': 'array'},
'type': {'const': 'mixed', 'title': 'Type'},
},
'required': ['type', 'sub_models'],
'title': 'MyModel',
'type': 'object',
}

0 comments on commit 8898b8f

Please sign in to comment.