Fix nested discriminated union schema gen, pt 2

[sydney-runkle]

Trying to keep this change as simple as possible -- more _generate_schema.py refactoring to come in future PRs.

This removes the usage of metadata keys to manage tagged union application in the latter part of schema simplification. I've made a comment on the most important line that fixed the bug we were experiencing with recursive discriminated unions.

See the code snippet below as an example of:
a) vastly increased performance compared to the faulty 2.6 logic
b) correct schema generation

from __future__ import annotations
import time
from typing import Literal, Annotated

from pydantic import Field, TypeAdapter, BaseModel
from pydantic._internal._core_utils import pretty_print_core_schema

class NestedState(BaseModel):
    state_type: Literal["nested"]
    substate: AnyState

class LoopState(BaseModel):
    state_type: Literal["loop"]
    substate: AnyState

class LeafState(BaseModel):
    state_type: Literal["leaf"]

AnyState = Annotated[NestedState | LoopState | LeafState, Field(..., discriminator="state_type")]

def build_nested_state(n):
    if n <= 0:
        return {"state_type": "leaf"}
    else:
        return {"state_type": "loop", "substate": {"state_type": "nested", "substate": build_nested_state(n-1)}}
        
adapter = TypeAdapter(AnyState)

start = time.time()
adapter.validate_python(build_nested_state(9))
print(time.time() - start)
#> 6.818771362304688e-05

start = time.time()
adapter.validate_python(build_nested_state(10))
print(time.time() - start)
#> 1.6927719116210938e-05

start = time.time()
adapter.validate_python(build_nested_state(11))
print(time.time() - start)
#> 1.6927719116210938e-05

pretty_print_core_schema(adapter.core_schema, True)
"""
{
    'type': 'definitions',
    'schema': {
        'type': 'tagged-union',
        'choices': {
            'nested': {'type': 'definition-ref', 'schema_ref': '__main__.NestedState:4629701792'},
            'loop': {'type': 'definition-ref', 'schema_ref': '__main__.LoopState:4898226368'},
            'leaf': {'type': 'definition-ref', 'schema_ref': '__main__.LeafState:4898242800'}
        },
        'discriminator': 'state_type',
        'strict': False,
        'from_attributes': True
    },
    'definitions': [
        {
            'type': 'model',
            'cls': <class '__main__.LeafState'>,
            'schema': {'type': 'model-fields', 'fields': {'state_type': {'type': 'model-field', 'schema': {'type': 'literal', 'expected': ['leaf']}}}, 'model_name': 'LeafState'},
            'ref': '__main__.LeafState:4898242800'
        },
        {
            'type': 'model',
            'cls': <class '__main__.LoopState'>,
            'schema': {
                'type': 'model-fields',
                'fields': {
                    'state_type': {'type': 'model-field', 'schema': {'type': 'literal', 'expected': ['loop']}},
                    'substate': {
                        'type': 'model-field',
                        'schema': {
                            'type': 'default',
                            'schema': {
                                'type': 'tagged-union',
                                'choices': {
                                    'nested': {'type': 'definition-ref', 'schema_ref': '__main__.NestedState:4629701792'},
                                    'loop': {'type': 'definition-ref', 'schema_ref': '__main__.LoopState:4898226368'},
                                    'leaf': {'type': 'definition-ref', 'schema_ref': '__main__.LeafState:4898242800'}
                                },
                                'discriminator': 'state_type',
                                'strict': False,
                                'from_attributes': True
                            },
                            'default': Ellipsis
                        }
                    }
                },
                'model_name': 'LoopState'
            },
            'ref': '__main__.LoopState:4898226368'
        },
        {
            'type': 'model',
            'cls': <class '__main__.NestedState'>,
            'schema': {
                'type': 'model-fields',
                'fields': {
                    'state_type': {'type': 'model-field', 'schema': {'type': 'literal', 'expected': ['nested']}},
                    'substate': {
                        'type': 'model-field',
                        'schema': {
                            'type': 'default',
                            'schema': {
                                'type': 'tagged-union',
                                'choices': {
                                    'nested': {'type': 'definition-ref', 'schema_ref': '__main__.NestedState:4629701792'},
                                    'loop': {'type': 'definition-ref', 'schema_ref': '__main__.LoopState:4898226368'},
                                    'leaf': {'type': 'definition-ref', 'schema_ref': '__main__.LeafState:4898242800'}
                                },
                                'discriminator': 'state_type',
                                'strict': False,
                                'from_attributes': True
                            },
                            'default': Ellipsis
                        }
                    }
                },
                'model_name': 'NestedState'
            },
            'ref': '__main__.NestedState:4629701792'
        }
    ]
}
"""

The above schema is now actually correct, and isn't cluttered with metadata tags relating to the application of discriminated unions :)

[cloudflare-pages]

Deploying with    Cloudflare Pages

Latest commit:	2ad301b
Status:	✅  Deploy successful!
Preview URL:	https://f3ed5104.pydantic-docs2.pages.dev
Branch Preview URL:	https://fix-discriminators-for-good.pydantic-docs2.pages.dev

View logs

[codspeed-hq]

CodSpeed Performance Report

Merging #8932 will not alter performance

_{Comparing fix-discriminators-for-good (2ad301b) with main (4fed81b)}

Summary

✅ 10 untouched benchmarks

[sydney-runkle]

Please review

[sydney-runkle]

Actually get rid of this once we use it

[sydney-runkle]

Important!!! If a choice is of type definition-ref, we want to just reuse that ref for the given choice. Before, we were going through this whole thing of fetching the value from definitions, then using that, but that ends up not working for nested / recursive schemas.

Our schema walking logic walks through both the schema and the definitions, so we can rest easy knowing that unions will be converted to tagged unions in the definitions list as well.

[dmontagu]

Does this work with JSON schema generation if e.g. the ref'ed schema is itself a discriminated union? Maybe that can't happen, and either way this seems like an improvement if no tests fail, but still

[sydney-runkle]

Good question. I believe that the walk core schema logic handles definitions schemas such that the function being applied during the walk is also applied to all of the definitions in a definitions schema, so that's why I felt comfortable removing this step.

This can be seen via the example test I added - discriminated union transformation logic is applied to the 2 schemas in the definitions list that require said changes!

[sydney-runkle]

I'll note, we do trace all of the way down each schema now to check for discriminated unions, which isn't the most performant, although we certainly get some nice benefits here from using refs instead of constantly pulling from definitions and then having to simplify nasty schemas.

I think it'd be beneficial for us to have some codecov tests for schema building, etc for PRs like this one.

[sydney-runkle]

If you expand this section, you can see that this test showcases the example I've shown in the PR description 👍

[adriangb]

Great work! The edited test shows that a model rebuild is no longer necessary, and hence this does fix things.

[sydney-runkle]

I've added a few tests showing the now-correct behavior for the example given in #7487, along with some extra variants for BaseModel and the pydantic dataclass.

Though these tests could be consolidated a bit with parametrization, I think that each of these test cases is simple enough that it's fine to leave them as singular tests that can easily be run independently.

[sydney-runkle]

Regarding my comment about other high-priority fixes for discriminated unions, this was one I had in mind: #8628.

I think we could include that fix + this improvement in 2.6.4.