Apr 2, 2025 · Apr 1, 2025
diff --git a/pydantic/_internal/_generate_schema.py b/pydantic/_internal/_generate_schema.py
            return _discriminated_union.apply_discriminator(
                schema,
                discriminator,
                self.defs._definitions,
            )
        except _discriminated_union.MissingDefinitionForUnionRef:
            # defer until defs are resolved

        remaining_defs: dict[str, CoreSchema] = {}

        # Note: this logic doesn't play well when core schemas with deferred discriminator metadata
        # and references are encountered. See the `test_deferred_discriminated_union_and_references()` test.
        for ref, inlinable_def_ref in gather_result['collected_references'].items():
            if inlinable_def_ref is not None and (inlining_behavior := _inlining_behavior(inlinable_def_ref)) != 'keep':
                if inlining_behavior == 'inline':
diff --git a/tests/test_discriminated_union.py b/tests/test_discriminated_union.py
 import pytest
 from dirty_equals import HasRepr, IsStr
 from pydantic_core import SchemaValidator, core_schema
 from typing_extensions import TypedDict
 from typing_extensions importTypeAliasType,TypedDict

 from pydantic import (
    BaseModel,
    ValidationError,
    field_validator,
 )
 from pydantic._internal._config import ConfigWrapper
 from pydantic._internal._discriminated_union import apply_discriminator
 from pydantic._internal._generate_schema import GenerateSchema
 from pydantic.dataclasses import dataclass as pydantic_dataclass
 from pydantic.errors import PydanticUserError
 from pydantic.fields import FieldInfo
    # `'pydantic_internal_union_discriminator'` meta key from the schemas. The schema cleaning
    # logic of `Reference` would then re-apply the deferred discriminator logic for `Base`.
    assert Base.__pydantic_core_schema__ == base_schema


 def test_discriminated_union_type_alias_type() -> None:
    """https://github.com/pydantic/pydantic/issues/11661

    This was fixed by making sure we provide the available definitions
    when first trying to apply discriminated unions during core schema
    generation (which we forgot to do). Our schema cleaning logic is still
    not working correctly when deferred discriminated unions are involved
    together with referenceable core schemas that should be inlined. In practice,
    I don't know if such a scenario can happen (see the test below --
    `test_deferred_discriminated_union_and_references()` for a theoretical example).
    """

    class Foo(BaseModel):
        type: Literal['foo'] = 'foo'

    Disc = TypeAliasType('Disc', Annotated[Foo, Field(discriminator='type')])

    class Main(BaseModel):
        f: Disc

    # Use the JSON Schema to avoid making assertions on the core schema, that
    # may be less stable:
    assert Main.model_json_schema()['$defs']['Disc']['discriminator'] == {
        'mapping': {'foo': '#/$defs/Foo'},
        'propertyName': 'type',
    }


 @pytest.mark.xfail(reason='deferred discriminated union info is lost on core schemas that are inlined.')
 def test_deferred_discriminated_union_and_references() -> None:
    class Foo(BaseModel):
        type: Literal['foo'] = 'foo'

    class Bar(BaseModel):
        type: Literal['bar'] = 'bar'

    gen_schema = GenerateSchema(ConfigWrapper(None))

    foo_ref = gen_schema.defs.create_definition_reference_schema(Foo.__pydantic_core_schema__)
    bar_ref = gen_schema.defs.create_definition_reference_schema(Bar.__pydantic_core_schema__)

    disc_union = core_schema.union_schema(
        choices=[foo_ref, bar_ref],
        metadata={'pydantic_internal_union_discriminator': 'type'},
        ref='disc_union',
    )
    disc_union_ref = gen_schema.defs.create_definition_reference_schema(disc_union)

    final_schema = gen_schema.clean_schema(disc_union_ref)

    assert final_schema['type'] == 'tagged-union'
Original file line number	Diff line number	Diff line change
Expand Up		@@ -644,6 +644,7 @@ def _apply_discriminator_to_union(
		return _discriminated_union.apply_discriminator(
		schema,
		discriminator,
		self.defs._definitions,
		)
		except _discriminated_union.MissingDefinitionForUnionRef:
		# defer until defs are resolved
Expand DownExpand Up		@@ -2766,6 +2767,8 @@ def finalize_schema(self, schema: CoreSchema) -> CoreSchema:

		remaining_defs: dict[str, CoreSchema] = {}

		# Note: this logic doesn't play well when core schemas with deferred discriminator metadata
		# and references are encountered. See the `test_deferred_discriminated_union_and_references()` test.
		for ref, inlinable_def_ref in gather_result['collected_references'].items():
		if inlinable_def_ref is not None and (inlining_behavior := _inlining_behavior(inlinable_def_ref)) != 'keep':
		if inlining_behavior == 'inline':
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -10,7 +10,7 @@
		import pytest
		from dirty_equals import HasRepr, IsStr
		from pydantic_core import SchemaValidator, core_schema
		from typing_extensions import TypedDict
		from typing_extensions importTypeAliasType,TypedDict

		from pydantic import (
		BaseModel,
Expand All		@@ -23,7 +23,9 @@
		ValidationError,
		field_validator,
		)
		from pydantic._internal._config import ConfigWrapper
		from pydantic._internal._discriminated_union import apply_discriminator
		from pydantic._internal._generate_schema import GenerateSchema
		from pydantic.dataclasses import dataclass as pydantic_dataclass
		from pydantic.errors import PydanticUserError
		from pydantic.fields import FieldInfo
Expand DownExpand Up		@@ -2199,3 +2201,56 @@ class Reference(BaseModel):
		# `'pydantic_internal_union_discriminator'` meta key from the schemas. The schema cleaning
		# logic of `Reference` would then re-apply the deferred discriminator logic for `Base`.
		assert Base.__pydantic_core_schema__ == base_schema


		def test_discriminated_union_type_alias_type() -> None:
		"""https://github.com/pydantic/pydantic/issues/11661

		This was fixed by making sure we provide the available definitions
		when first trying to apply discriminated unions during core schema
		generation (which we forgot to do). Our schema cleaning logic is still
		not working correctly when deferred discriminated unions are involved
		together with referenceable core schemas that should be inlined. In practice,
		I don't know if such a scenario can happen (see the test below --
		`test_deferred_discriminated_union_and_references()` for a theoretical example).
		"""

		class Foo(BaseModel):
		type: Literal['foo'] = 'foo'

		Disc = TypeAliasType('Disc', Annotated[Foo, Field(discriminator='type')])

		class Main(BaseModel):
		f: Disc

		# Use the JSON Schema to avoid making assertions on the core schema, that
		# may be less stable:
		assert Main.model_json_schema()['$defs']['Disc']['discriminator'] == {
		'mapping': {'foo': '#/$defs/Foo'},
		'propertyName': 'type',
		}


		@pytest.mark.xfail(reason='deferred discriminated union info is lost on core schemas that are inlined.')
		def test_deferred_discriminated_union_and_references() -> None:
		class Foo(BaseModel):
		type: Literal['foo'] = 'foo'

		class Bar(BaseModel):
		type: Literal['bar'] = 'bar'

		gen_schema = GenerateSchema(ConfigWrapper(None))

		foo_ref = gen_schema.defs.create_definition_reference_schema(Foo.__pydantic_core_schema__)
		bar_ref = gen_schema.defs.create_definition_reference_schema(Bar.__pydantic_core_schema__)

		disc_union = core_schema.union_schema(
		choices=[foo_ref, bar_ref],
		metadata={'pydantic_internal_union_discriminator': 'type'},
		ref='disc_union',
		)
		disc_union_ref = gen_schema.defs.create_definition_reference_schema(disc_union)

		final_schema = gen_schema.clean_schema(disc_union_ref)

		assert final_schema['type'] == 'tagged-union'