Fix regression in core schema generation for indirect definition references (#8702)

dmontagu · sydney-runkle · sydney-runkle · commit aac1069ec33d · 2024-02-05T07:32:44.000-06:00
Co-authored-by: Sydney Runkle &lt;54324534+sydney-runkle@users.noreply.github.com&gt;
diff --git 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
@@ -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',
+    }
