1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
|
from __future__ import annotations
from typing import Annotated, Literal, Union
import pytest
from pydantic import BaseModel, Field, TypeAdapter
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[Union[NestedState, LoopState, LeafState], Field(discriminator='state_type')]
@pytest.mark.benchmark
def test_schema_build(benchmark) -> None:
@benchmark
def run():
adapter = TypeAdapter(AnyState)
assert adapter.core_schema['schema']['type'] == 'tagged-union'
any_state_adapter = TypeAdapter(AnyState)
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)}}
@pytest.mark.benchmark
def test_efficiency_with_highly_nested_examples(benchmark) -> None:
# can go much higher, but we keep it reasonably low here for a proof of concept
@benchmark
def run():
for i in range(1, 12):
very_nested_input = build_nested_state(i)
any_state_adapter.validate_python(very_nested_input)
|
