Fix discriminated union schema gen bug #8904

sydney-runkle · 2024-02-27T16:30:19Z

With this code snippet:

from __future__ import annotations
from typing import Literal, Annotated
from pydantic import Field, TypeAdapter, BaseModel

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

# If this type is left out, the model behaves normally again
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)

Previously:

# the next statement takes around 0.8s
adapter.validate_python(build_nested_state(9))

# the next statement takes around 3.5s
adapter.validate_python(build_nested_state(10))

# the next statement takes around 12.5s
adapter.validate_python(build_nested_state(11))

Now, note the speedups:

start = time.time()
adapter.validate_python(build_nested_state(9))
print(time.time() - start)
#> 0.0063190460205078125

start = time.time()
adapter.validate_python(build_nested_state(10))
print(time.time() - start)
#> 0.02297687530517578

start = time.time()
adapter.validate_python(build_nested_state(11))
print(time.time() - start)
#> 0.057311058044433594

codspeed-hq · 2024-02-27T16:33:28Z

CodSpeed Performance Report

Merging #8904 will not alter performance

_{Comparing fix-schema-building-bug (d35cd68) with main (752daaf)}

Summary

✅ 10 untouched benchmarks

cloudflare-pages · 2024-02-27T16:54:04Z

Deploying with Cloudflare Pages

Latest commit:	`d35cd68`
Status:	✅ Deploy successful!
Preview URL:	https://36f1154d.pydantic-docs2.pages.dev
Branch Preview URL:	https://fix-schema-building-bug.pydantic-docs2.pages.dev

View logs

adriangb

I think I'd like a little bit more documentation on how this fixes things (maybe resultant in a comment in the code) so we don't make the same mistake again, but other than that LGTM.

@sydney-runkle this is really impressive work in record time, great job.

adriangb · 2024-02-27T17:27:32Z

pydantic/_internal/_discriminated_union.py

@@ -54,6 +54,7 @@ def inner(s: core_schema.CoreSchema, recurse: _core_utils.Recurse) -> core_schem
        if discriminator is not None:
            if definitions is None:
                definitions = collect_definitions(schema)
+                definitions = {k: recurse(v, inner) for k, v in definitions.items()}
            s = apply_discriminator(s, discriminator, definitions)


🤯 care to explain how this fixes things?

Sure thing! I'll add some more comments in the code to explain things 😄!

I'll also add a more in-depth note to this PR, so we can reference that for more detail, if needed 👍.

So previously, with the above snippet of code, the resultant core_schema was the following:

{'definitions': [{'cls': <class '__main__.LeafState'>, 'config': {'title': 'LeafState'}, 'custom_init': False, 'metadata': {'pydantic.internal.needs_apply_discriminated_union': False, 'pydantic_js_annotation_functions': [], 'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x101bcd900>, cls=<class '__main__.LeafState'>), <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.LeafState'>>]}, 'ref': '__main__.LeafState:5484303040', 'root_model': False, 'schema': {'computed_fields': [], 'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101c09d80>], 'pydantic_js_functions': []}, 'schema': {'expected': ['leaf'], 'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'type': 'literal'}, 'type': 'model-field'}}, 'model_name': 'LeafState', 'type': 'model-fields'}, 'type': 'model'}, {'cls': <class '__main__.LoopState'>, 'config': {'title': 'LoopState'}, 'custom_init': False, 'metadata': {'pydantic.internal.needs_apply_discriminated_union': False, 'pydantic_js_annotation_functions': [], 'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x101bcd900>, cls=<class '__main__.LoopState'>), <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.LoopState'>>]}, 'ref': '__main__.LoopState:5484288784', 'root_model': False, 'schema': {'computed_fields': [], 'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101cf8310>], 'pydantic_js_functions': []}, 'schema': {'expected': ['loop'], 'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'type': 'literal'}, 'type': 'model-field'}, 'substate': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101cf83a0>], 'pydantic_js_functions': []}, 'schema': {'default': Ellipsis, 'schema': {'choices': [{'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'schema_ref': '__main__.NestedState:5484231952', 'type': 'definition-ref'}, {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'schema_ref': '__main__.LoopState:5484288784', 'type': 'definition-ref'}, {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'schema_ref': '__main__.LeafState:5484303040', 'type': 'definition-ref'}], 'metadata': {'pydantic.internal.needs_apply_discriminated_union': True, 'pydantic.internal.union_discriminator': 'state_type'}, 'type': 'union'}, 'type': 'default'}, 'type': 'model-field'}}, 'model_name': 'LoopState', 'type': 'model-fields'}, 'type': 'model'}, {'cls': <class '__main__.NestedState'>, 'config': {'title': 'NestedState'}, 'custom_init': False, 'metadata': {'pydantic.internal.needs_apply_discriminated_union': False, 'pydantic_js_annotation_functions': [], 'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x101bcd900>, cls=<class '__main__.NestedState'>), <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.NestedState'>>]}, 'ref': '__main__.NestedState:5484231952', 'root_model': False, 'schema': {'computed_fields': [], 'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101cf80d0>], 'pydantic_js_functions': []}, 'schema': {'expected': ['nested'], 'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'type': 'literal'}, 'type': 'model-field'}, 'substate': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101cf8160>], 'pydantic_js_functions': []}, 'schema': {'default': Ellipsis, 'schema': {'choices': [{'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'schema_ref': '__main__.NestedState:5484231952', 'type': 'definition-ref'}, {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'schema_ref': '__main__.LoopState:5484288784', 'type': 'definition-ref'}, {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False}, 'schema_ref': '__main__.LeafState:5484303040', 'type': 'definition-ref'}], 'metadata': {'pydantic.internal.needs_apply_discriminated_union': True, 'pydantic.internal.union_discriminator': 'state_type'}, 'type': 'union'}, 'type': 'default'}, 'type': 'model-field'}}, 'model_name': 'NestedState', 'type': 'model-fields'}, 'type': 'model'}], 'schema': {'choices': {'leaf': {'schema_ref': '__main__.LeafState:5484303040', 'type': 'definition-ref'}, 'loop': {'schema_ref': '__main__.LoopState:5484288784', 'type': 'definition-ref'}, 'nested': {'schema_ref': '__main__.NestedState:5484231952', 'type': 'definition-ref'}}, 'discriminator': 'state_type', 'from_attributes': True, 'metadata': {'pydantic.internal.needs_apply_discriminated_union': True, 'pydantic.internal.union_discriminator': 'state_type'}, 'strict': False, 'type': 'tagged-union'}, 'type': 'definitions'}

We already identified that the schemas for NestedState and LoopState within the definitions list were marked as having union type substate schemas, rather than tagged-union substate schemas. Interestingly, though, the metadata for those schemas had the necessary discriminated-union related information:

'metadata': {'pydantic.internal.needs_apply_discriminated_union': True, 'pydantic.internal.union_discriminator': 'state_type'}

So the question, then, is why aren't we "applying" the discriminators to those schemas? This is where my change comes in - after we collect the definitions schemas, we also need to pass over them + apply discriminators where appropriate. Within the inner function, we were already calling recurse, but we weren't calling it on any definitions schemas, which was the missing piece.

Totally makes sense!

samuelcolvin

LGTM!

sydney-runkle · 2024-03-01T22:12:45Z

Update -- I believe this fixed the performance issue for the above case, but didn't entirely fix the schema generation process. Here's a fix for that: #8932.

I'd like to work on a few other schema generation / performance related issues, but hopefully we can get these out in another patch release soon 😄.

Feel free to ping me if you have any questions!

fixgit statusgit status

517c038

sydney-runkle added the relnotes-fix Used for bugfixes. label Feb 27, 2024

add test

44ec9a8

fix for earlier python versions

3b25873

sydney-runkle requested review from adriangb and davidhewitt February 27, 2024 17:10

adriangb approved these changes Feb 27, 2024

View reviewed changes

add note about reasoning behind fix

d35cd68

adriangb approved these changes Feb 27, 2024

View reviewed changes

samuelcolvin approved these changes Feb 27, 2024

View reviewed changes

sydney-runkle merged commit 5a4c056 into main Feb 27, 2024
54 checks passed

sydney-runkle deleted the fix-schema-building-bug branch February 27, 2024 18:05

sydney-runkle added a commit that referenced this pull request Feb 27, 2024

Fix discriminated union schema gen bug (#8904)

6fcee87

sydney-runkle mentioned this pull request Feb 27, 2024

Nested/recursive Discriminated unions don't seem to work #7487

Closed

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix discriminated union schema gen bug #8904

Fix discriminated union schema gen bug #8904

sydney-runkle commented Feb 27, 2024

codspeed-hq bot commented Feb 27, 2024 •

edited

cloudflare-pages bot commented Feb 27, 2024 •

edited

adriangb left a comment

adriangb Feb 27, 2024

sydney-runkle Feb 27, 2024 •

edited

sydney-runkle Feb 27, 2024

adriangb Feb 27, 2024

adriangb Feb 27, 2024

samuelcolvin left a comment

sydney-runkle commented Mar 1, 2024

Fix discriminated union schema gen bug #8904

Fix discriminated union schema gen bug #8904

Conversation

sydney-runkle commented Feb 27, 2024

codspeed-hq bot commented Feb 27, 2024 • edited

CodSpeed Performance Report

Merging #8904 will not alter performance

Summary

cloudflare-pages bot commented Feb 27, 2024 • edited

Deploying with Cloudflare Pages

adriangb left a comment

Choose a reason for hiding this comment

adriangb Feb 27, 2024

Choose a reason for hiding this comment

sydney-runkle Feb 27, 2024 • edited

Choose a reason for hiding this comment

sydney-runkle Feb 27, 2024

Choose a reason for hiding this comment

adriangb Feb 27, 2024

Choose a reason for hiding this comment

adriangb Feb 27, 2024

Choose a reason for hiding this comment

samuelcolvin left a comment

Choose a reason for hiding this comment

sydney-runkle commented Mar 1, 2024

codspeed-hq bot commented Feb 27, 2024 •

edited

cloudflare-pages bot commented Feb 27, 2024 •

edited

sydney-runkle Feb 27, 2024 •

edited