Fixes discriminated union bug regression when using enums

[kfreezen]

Change Summary

Fixes regression documented in pydantic/pydantic#9235, which was introduced in upgrade to pyo3 v0.21

Related issue number

pydantic/pydantic#9235

Checklist

• [ ] Unit tests for the changes exist
• [ ] Documentation reflects the changes where applicable
• [ ] Pydantic tests pass with this pydantic-core (except for expected changes)
• [x] My PR is ready to review, please add a comment including the phrase "please review" to assign reviewers

Selected Reviewer: @dmontagu

[kfreezen]

@sydney-runkle should I create a PR removing the xfail in pydantic as well? How do you handle upstream changes like this?

[kfreezen]

please review

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

:loudspeaker: Thoughts on this report? Let us know!

[codspeed-hq[bot]]

CodSpeed Performance Report

Merging #1286 will not alter performance

_{Comparing kfreezen:discriminated-union-bug (1e69121) with main (0e6b377)}

Summary

✅ 155 untouched benchmarks

[kfreezen]

@dmontagu you know what, I think I got ahead of myself here, I need to finagle a test for this change.

[quizmoon]

Would be great to have it fixed as soon as possible, is there any prediction date / version? Thank you.

[sydney-runkle]

@kfreezen,

Amazing work here!

Here's a test we can add to the tests/serializers/test_literal.py file:

def test_literal_with_enum() -> None:
    class SomeEnum(str, Enum):
        CAT = 'cat'
        DOG = 'dog'

    @dataclass
    class Dog:
        name: str
        type: Literal[SomeEnum.DOG] = SomeEnum.DOG

    @dataclass
    class Cat:
        name: str
        type: Literal[SomeEnum.CAT] = SomeEnum.CAT

    @dataclass
    class Yard:
        pet: Union[Dog, Cat]

    serializer = SchemaSerializer(
        core_schema.model_schema(
            cls=Yard, 
            schema=core_schema.model_fields_schema(
                fields={
                    'pet': core_schema.model_field(
                        schema=core_schema.tagged_union_schema(
                            choices={
                                SomeEnum.DOG: core_schema.model_schema(
                                    cls=Dog,
                                    schema=core_schema.model_fields_schema(
                                        fields={
                                            'type': core_schema.model_field(
                                                schema=core_schema.with_default_schema(
                                                    schema=core_schema.literal_schema([SomeEnum.DOG]),
                                                    default=SomeEnum.DOG
                                                )
                                            ),
                                            'name': core_schema.model_field(schema=core_schema.str_schema())
                                        },
                                        model_name='Dog'
                                    )
                                ),
                                SomeEnum.CAT: core_schema.model_schema(
                                    cls=Cat,
                                    schema=core_schema.model_fields_schema(
                                        fields={
                                            'type': core_schema.model_field(
                                                schema=core_schema.with_default_schema(
                                                    schema=core_schema.literal_schema([SomeEnum.CAT]),
                                                    default=SomeEnum.CAT
                                                )
                                            ),
                                            'name': core_schema.model_field(schema=core_schema.str_schema())
                                        },
                                        model_name='Cat'
                                    )
                                )
                            },
                            discriminator='type',
                            strict=False,
                            from_attributes=True
                        )
                    )
                }
            )
        )
    )

    yard = Yard(pet=Dog(name='Rex'))
    assert serializer.to_python(yard, mode='json') == {'pet': {'type': 'dog', 'name': 'Rex'}}

[sydney-runkle]

And yep, when we bring this into pydantic, we'll remove that xfail.

[sydney-runkle]

Once you add that test, I'll approve and we can merge :). Maybe give it a more specific name related to the union, just for clarity.

[kfreezen]

I’ll get that added tomorrow night

On Tue, Jul 2, 2024 at 20:39 Sydney Runkle @.***> wrote:

Once you add that test, I'll approve and we can merge :)

— Reply to this email directly, view it on GitHub https://github.com/pydantic/pydantic-core/pull/1286#issuecomment-2204694685, or unsubscribe https://github.com/notifications/unsubscribe-auth/AACUWLHRXQFIWGMXP5JFYALZKM235AVCNFSM6AAAAABHCZHBVWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEMBUGY4TINRYGU . You are receiving this because you were mentioned.Message ID: @.***>

[kfreezen]

Whoops, sorry guys, I completely forgot about this. Been on a sabbatical for the last 4 weeks. I will add it right now.

[quizmoon]

@kfreezen great job!

@sydney-runkle could you please go forward with it?