diff --git a/pydantic/_internal/_core_utils.py b/pydantic/_internal/_core_utils.py index 2c2f027af4..432de5f565 100644 --- a/pydantic/_internal/_core_utils.py +++ b/pydantic/_internal/_core_utils.py @@ -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) diff --git a/tests/test_discriminated_union.py b/tests/test_discriminated_union.py index 24e319da30..fb12e2d661 100644 --- a/tests/test_discriminated_union.py +++ b/tests/test_discriminated_union.py @@ -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', + }