From 81bde4d3e13560b8dfe5d2dcac6c54369a4ea38b Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 15:58:59 +0000 Subject: [PATCH 1/2] Compile nested dataclass schemas along with the outer one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A nested dataclass field was built as a plain Schema around All(inner, _Constructor). A plain Schema sees only the All, so it never armed under any policy, and DataclassSchema.compile() generated the outer mapping alone: every level below it stayed interpreted, which is why compiling a dataclass tree barely moved its numbers. The builder now makes a fused dataclass Schema subclass for every level, carrying the compile logic DataclassSchema had. It arms and generates like any schema, and compiling it compiles the nested dataclass schemas under its fields first (through a list, a dict, Maybe, All, or a union), so a tree generates all the way down. The outer compile flag is threaded into every nested schema, so compile=False keeps the whole tree interpreted. A Schema the user placed in an Annotated keeps its own flag and is not entered. A second compile() is a no-op instead of stacking validators. A schema on a recursion cycle must not compile: a generated validator bails to its interpreted twin on a failure, and the twin re-enters the generated code below it, so every compiled level on a cycle doubles a failing validation once per turn, exponential in the data's depth. Taking a recursive reference now marks every schema from its target down to the field taking it, and a marked schema never arms or generates; the schema above the cycle still does. A regression test pins the deep failure to linear time. On a three-level API response the compiled path drops from about 22µs to about 10µs per document. A nested-dataclass CodSpeed pair tracks it. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ANcsDvtNPxCRTk33uagPak --- bench/test_benchmarks.py | 46 +++ .../content/docs/guides/compiled-schemas.md | 10 +- src/probatio/dataclass_schema.py | 276 ++++++++++++++---- src/probatio/decorator.py | 8 +- tests/__snapshots__/test_public_surface.ambr | 2 + tests/test_codegen.py | 199 ++++++++++++- 6 files changed, 477 insertions(+), 64 deletions(-) diff --git a/bench/test_benchmarks.py b/bench/test_benchmarks.py index 5cef6f02..c5cdb588 100644 --- a/bench/test_benchmarks.py +++ b/bench/test_benchmarks.py @@ -367,12 +367,58 @@ def test_validate_dataclass_compiled(benchmark: Any) -> None: assert result.port == 443 +@dataclass +class _Reading: + """A unit/value pair, the nested leaf of a dataclass tree.""" + + unit: str + value: float + + +@dataclass +class _Device: + """A dataclass holding nested dataclasses: the shape an API response has. + + Compiling the outer dataclass compiles the nested schemas with it, so this is + where a regression in that recursion shows up; the flat ``_Service`` above + cannot see one. + """ + + name: str + power: _Reading + voltage: _Reading + current: _Reading | None = None + + +DEVICE = probatio.DataclassSchema(_Device, compile=False) +DEVICE_COMPILED = probatio.DataclassSchema(_Device).compile() +DEVICE_PAYLOAD = { + "name": "inverter", + "power": {"unit": "W", "value": 300.0}, + "voltage": {"unit": "V", "value": 230.1}, + "current": {"unit": "A", "value": 1.3}, +} + + +def test_validate_dataclass_tree(benchmark: Any) -> None: + """Validate and construct a dataclass with nested dataclasses (interpreted).""" + result = benchmark(DEVICE, DEVICE_PAYLOAD) + assert result.power.value == 300.0 + + +def test_validate_dataclass_tree_compiled(benchmark: Any) -> None: + """Validate and construct a dataclass tree, every level through generated code.""" + result = benchmark(DEVICE_COMPILED, DEVICE_PAYLOAD) + assert result.power.value == 300.0 + + def test_compiled_benchmarks_use_the_generated_validator() -> None: """Guard: every ``*_compiled`` benchmark really measures generated code.""" assert _is_compiled(CONFIG_COMPILED) assert _is_compiled(LEAF_HEAVY_COMPILED) assert _is_compiled(NESTED_COMPILED) assert _is_compiled(SERVICE_COMPILED) + assert _is_compiled(DEVICE_COMPILED) assert _is_compiled(LIST_COMPILED) diff --git a/docs/src/content/docs/guides/compiled-schemas.md b/docs/src/content/docs/guides/compiled-schemas.md index 50935759..97f13091 100644 --- a/docs/src/content/docs/guides/compiled-schemas.md +++ b/docs/src/content/docs/guides/compiled-schemas.md @@ -63,7 +63,10 @@ schema = Schema({Required("name"): str, Optional("age"): int}).compile() `DataclassSchema` and `TypedDictSchema` take the same flag, and they are the clearest win: compiling fuses field validation and object construction into one -function with no intermediate dict. +function with no intermediate dict. A nested dataclass field compiles along with +the dataclass that holds it, so a tree of dataclasses (an API response, say) +generates all the way down, and the flag covers the whole tree: `compile=False` +on the outer dataclass keeps every nested one interpreted too. ```python from dataclasses import dataclass @@ -123,8 +126,9 @@ inlines the validators that show up most: `Coerce` to `int`/`float`, numeric (`Any("on", "off")`, validated as a membership test), `Maybe(X)` (an optional validated value), and `All`/`Any` chains of those. -Anything it does not handle is simply left interpreted. A recursive schema (`Self` -or a recursive JSON Schema `$ref`), an exclusive or inclusive group, a multi-element +Anything it does not handle is simply left interpreted. A recursive schema (`Self`, +a recursive JSON Schema `$ref`, or a dataclass on a recursion cycle: one that holds +itself, or two that hold each other), an exclusive or inclusive group, a multi-element list (an "each item matches one of these" union), a list whose element does not inline, an exotic key: the schema works exactly as before, just without the speedup. You never lose correctness by compiling; at worst you do not gain speed. diff --git a/src/probatio/dataclass_schema.py b/src/probatio/dataclass_schema.py index 7e8a4212..403f5384 100644 --- a/src/probatio/dataclass_schema.py +++ b/src/probatio/dataclass_schema.py @@ -40,6 +40,7 @@ from collections.abc import ( Callable, Iterable, + Iterator, Mapping, MutableMapping, MutableSequence, @@ -113,11 +114,15 @@ class _RecursiveSchemaRef: ``Invalid`` rather than overflowing the stack, exactly like ``Self``. """ - __slots__ = ("_validate",) + __slots__ = ("_validate", "on_cycle") def __init__(self) -> None: """Start unbound; ``bind`` fills in the compiled schema before any call.""" self._validate: CompiledSchema | None = None + # Set once a field takes this reference (or one to an ancestor through + # it): the schema it stands for is on a recursion cycle and must stay + # interpreted. See ``_mark_cycle``. + self.on_cycle = False def bind(self, schema: Schema) -> None: """Point the reference at the finished schema's compiled engine.""" @@ -208,6 +213,8 @@ def _annotation_to_schema( # noqa: PLR0911, PLR0912 annotation: TypingAny, self_refs: dict[type, _RecursiveSchemaRef], extra: int, + *, + compiled: bool | None, ) -> TypingAny: """Map a single type annotation to a probatio schema fragment. @@ -220,6 +227,11 @@ def _annotation_to_schema( # noqa: PLR0911, PLR0912 nested dataclass, TypedDict, and container so the policy applies all the way down. A ``Key(extra=...)`` in a field's ``Annotated`` metadata overrides it for that field's subtree. + + ``compiled`` is the enclosing schema's ``compile`` flag, threaded into every + nested dataclass and TypedDict schema the same way, so ``compile=False`` on the + outer schema keeps the whole tree interpreted and ``compile=True`` generates + all the way down. """ if annotation is TypingAny: return object @@ -228,7 +240,9 @@ def _annotation_to_schema( # noqa: PLR0911, PLR0912 if hasattr(annotation, "__supertype__"): # A ``NewType`` is a thin alias; validate against the type it wraps. - return _annotation_to_schema(annotation.__supertype__, self_refs, extra) + return _annotation_to_schema( + annotation.__supertype__, self_refs, extra, compiled=compiled + ) if get_origin(annotation) is Annotated: # Callable Annotated metadata are extra validators; other metadata is for @@ -237,7 +251,9 @@ def _annotation_to_schema( # noqa: PLR0911, PLR0912 base, *meta = get_args(annotation) key = next((item for item in meta if isinstance(item, Key)), None) field_extra = key.extra if key is not None and key.extra is not None else extra - base_schema = _annotation_to_schema(base, self_refs, field_extra) + base_schema = _annotation_to_schema( + base, self_refs, field_extra, compiled=compiled + ) extras = [item for item in meta if callable(item)] if not extras: return base_schema @@ -263,15 +279,21 @@ def _annotation_to_schema( # noqa: PLR0911, PLR0912 if _is_dataclass_type(annotation): if annotation in self_refs: + _mark_cycle(self_refs, annotation) return self_refs[annotation] - return create_dataclass_schema(annotation, extra=extra, _self_refs=self_refs) + return create_dataclass_schema( + annotation, extra=extra, _self_refs=self_refs, _compile=compiled + ) if is_typeddict(annotation): # A nested TypedDict validates as its own mapping schema (it cannot be an # ``isinstance`` check, since a TypedDict has no runtime type). if annotation in self_refs: + _mark_cycle(self_refs, annotation) return self_refs[annotation] - return create_typeddict_schema(annotation, extra=extra, _self_refs=self_refs) + return create_typeddict_schema( + annotation, extra=extra, _self_refs=self_refs, _compile=compiled + ) origin = get_origin(annotation) if origin is None: @@ -281,7 +303,7 @@ def _annotation_to_schema( # noqa: PLR0911, PLR0912 args = get_args(annotation) if origin in (TypingUnion, types.UnionType): - return _union_to_schema(args, self_refs, extra) + return _union_to_schema(args, self_refs, extra, compiled=compiled) if origin is Literal: return In(list(args)) if not args: @@ -289,32 +311,32 @@ def _annotation_to_schema( # noqa: PLR0911, PLR0912 # to descend into, so validate the container type itself. return origin if origin is list: - return [_annotation_to_schema(args[0], self_refs, extra)] + return [_annotation_to_schema(args[0], self_refs, extra, compiled=compiled)] if origin is set: - return {_annotation_to_schema(args[0], self_refs, extra)} + return {_annotation_to_schema(args[0], self_refs, extra, compiled=compiled)} if origin is frozenset: - return frozenset([_annotation_to_schema(args[0], self_refs, extra)]) + return frozenset( + [_annotation_to_schema(args[0], self_refs, extra, compiled=compiled)] + ) if origin is dict: return { - _annotation_to_schema(args[0], self_refs, extra): _annotation_to_schema( - args[1], - self_refs, - extra, - ), + _annotation_to_schema( + args[0], self_refs, extra, compiled=compiled + ): _annotation_to_schema(args[1], self_refs, extra, compiled=compiled), } if origin is tuple: - return _tuple_to_schema(args, self_refs, extra) + return _tuple_to_schema(args, self_refs, extra, compiled=compiled) if origin in (Sequence, MutableSequence): # An abstract sequence validates element-wise like a list; config data # arrives as a list, and a str (also a Sequence) is correctly not one. - return [_annotation_to_schema(args[0], self_refs, extra)] + return [_annotation_to_schema(args[0], self_refs, extra, compiled=compiled)] if origin in (AbstractSet, MutableSet): - return {_annotation_to_schema(args[0], self_refs, extra)} + return {_annotation_to_schema(args[0], self_refs, extra, compiled=compiled)} if origin in (Mapping, MutableMapping): return { - _annotation_to_schema(args[0], self_refs, extra): _annotation_to_schema( - args[1], self_refs, extra - ), + _annotation_to_schema( + args[0], self_refs, extra, compiled=compiled + ): _annotation_to_schema(args[1], self_refs, extra, compiled=compiled), } # For other generics, validate the container type instead of accepting # everything. Some origins are not runtime types, so those stay open. @@ -325,6 +347,8 @@ def _union_to_schema( args: tuple[TypingAny, ...], self_refs: dict[type, _RecursiveSchemaRef], extra: int, + *, + compiled: bool | None, ) -> TypingAny: """Map a union: ``X | None`` becomes ``Maybe(X)``, a wider union becomes ``Any``. @@ -335,7 +359,10 @@ def _union_to_schema( """ has_none = type(None) in args member_types = [a for a in args if a is not type(None)] - members = [_annotation_to_schema(a, self_refs, extra) for a in member_types] + members = [ + _annotation_to_schema(a, self_refs, extra, compiled=compiled) + for a in member_types + ] if len(members) == 1: inner: TypingAny = members[0] @@ -414,10 +441,33 @@ def discriminant(value: TypingAny, validators: list[TypingAny]) -> list[TypingAn return discriminant +def _mark_cycle(self_refs: dict[type, _RecursiveSchemaRef], target: type) -> None: + """Flag every schema from ``target`` down to the current level as on a cycle. + + A recursive reference is only taken for a type still under construction, so + the schemas built between it and the field taking the reference all lie on + the cycle the reference closes (``self_refs`` is in build-stack order). None + of them may compile: a generated validator bails to its interpreted twin on a + failure, and the twin re-enters the generated code of the schemas below it, + so every compiled level on a cycle doubles the work of a failing validation + once per turn of the cycle, exponential in the data's depth. The schemas above + the cycle may compile: they are entered once, so they add one doubling, not + one per level of data. + """ + marking = False + for dataclass_type, ref in self_refs.items(): + if dataclass_type is target: + marking = True + if marking: + ref.on_cycle = True + + def _tuple_to_schema( args: tuple[TypingAny, ...], self_refs: dict[type, _RecursiveSchemaRef], extra: int, + *, + compiled: bool | None, ) -> TypingAny: """Map a tuple: ``tuple[X, ...]`` is homogeneous, ``tuple[X, Y]`` is positional. @@ -425,9 +475,11 @@ def _tuple_to_schema( and preserve whichever was given (the value is not coerced between them). """ if len(args) == 2 and args[1] is Ellipsis: - element = _annotation_to_schema(args[0], self_refs, extra) + element = _annotation_to_schema(args[0], self_refs, extra, compiled=compiled) return Any([element], (element,)) - return ExactSequence([_annotation_to_schema(a, self_refs, extra) for a in args]) + return ExactSequence( + [_annotation_to_schema(a, self_refs, extra, compiled=compiled) for a in args] + ) def _constant(value: TypingAny) -> TypingAny: @@ -648,6 +700,8 @@ def _field_mapping( constraints: dict[str, TypingAny], self_refs: dict[type, _RecursiveSchemaRef], extra: int, + *, + compiled: bool | None, ) -> dict[TypingAny, TypingAny]: """Build the mapping schema for a dataclass's constructor fields.""" try: @@ -666,7 +720,9 @@ def _field_mapping( mapping: dict[TypingAny, TypingAny] = {} for field in fields: annotation = _field_annotation(hints.get(field.name, TypingAny)) - value_schema = _annotation_to_schema(annotation, self_refs, extra) + value_schema = _annotation_to_schema( + annotation, self_refs, extra, compiled=compiled + ) if field.name in constraints: value_schema = All(value_schema, constraints[field.name]) @@ -764,6 +820,7 @@ def create_dataclass_schema( required: bool = False, extra: int = PREVENT_EXTRA, _self_refs: dict[type, _RecursiveSchemaRef] | None = None, + _compile: bool | None = None, ) -> Schema: """Build a ``Schema`` that validates a dict and constructs a dataclass. @@ -789,7 +846,11 @@ def create_dataclass_schema( ref = _RecursiveSchemaRef() self_refs[dataclass_type] = ref mapping = _field_mapping( - dataclass_type, additional_constraints or {}, self_refs, extra + dataclass_type, + additional_constraints or {}, + self_refs, + extra, + compiled=_compile, ) # DataclassSchema compiles mapping validation and construction together, so it @@ -808,7 +869,13 @@ def create_dataclass_schema( constructor = _Constructor( dataclass_type, remove, filter_keys=extra == ALLOW_EXTRA or bool(remove) ) - schema = Schema(All(inner, constructor)) + # A fused schema, not a plain one: it knows how to generate code for itself, + # so a nested dataclass compiles along with the schema that holds it (and on + # its own under the ``AUTO`` policy once it is hot), where a plain ``Schema`` + # sees only an ``All`` and never arms. + schema = _FusedDataclassSchema( + All(inner, constructor), on_cycle=ref.on_cycle, compile=_compile + ) # Fused before ``bind``, so a recursive field's captured engine is the fused # one too, not the slower combinator tower. _install_fused_interpreted(schema) @@ -896,7 +963,7 @@ def validate_splat(data: TypingAny) -> TypingAny: return validate_splat -def _install_fused_interpreted(schema: Schema) -> None: +def _install_fused_interpreted(schema: _FusedDataclassSchema) -> None: """Swap ``schema``'s interpreted engine for the fused validate-and-construct. ``schema`` wraps ``All(inner, _Constructor)``. The inner mapping schema is @@ -906,7 +973,9 @@ def _install_fused_interpreted(schema: Schema) -> None: The fused engine lands wherever the interpreted validator lives: in ``_interpreted`` when the schema armed for lazy compilation (it then also - becomes the generated code's bail-out fallback), else in ``_compiled``. + becomes the generated code's bail-out fallback), else in ``_compiled``. It is + also remembered as ``_fused_engine``, so ``_compile_from`` can tell a schema + that is still interpreted from one already generated. """ # Force a lazily-deferred wrapper to build first, so its arming and engine are # in place before the fused engine replaces the interpreted one. @@ -915,12 +984,39 @@ def _install_fused_interpreted(schema: Schema) -> None: if inner._uses_self: # noqa: SLF001 return fused = _fuse_validate_and_construct(inner._compiled, constructor) # noqa: SLF001 + schema._fused_engine = fused # noqa: SLF001 if "_interpreted" in schema.__dict__: schema._interpreted = fused # noqa: SLF001 else: schema._compiled = fused # noqa: SLF001 +def _built_dataclass_schemas(schema: TypingAny) -> Iterator[_FusedDataclassSchema]: + """Yield the nested dataclass schemas the builder made inside a field's schema. + + Walks through the containers and combinators a field annotation produces (a + list, a dict, ``Maybe``, ``All``, a union) to the dataclass schemas + ``create_dataclass_schema`` built for nested fields. Only those count: a + ``Schema`` the user placed in an ``Annotated`` (a ``DataclassSchema`` carrying + its own ``compile`` flag among them) decides for itself and is not entered. A + recursive reference has no children to enter, which is what keeps a cyclic + dataclass graph from walking forever. + """ + if type(schema) is _FusedDataclassSchema: + yield schema + elif isinstance(schema, (list, tuple, set, frozenset)): + for item in schema: + yield from _built_dataclass_schemas(item) + elif isinstance(schema, dict): + for value in schema.values(): + yield from _built_dataclass_schemas(value) + elif not isinstance(schema, Schema): + children = getattr(schema, "__probatio_child_schemas__", None) + if children is not None: + for child in children(): + yield from _built_dataclass_schemas(child) + + # Distinguishes "the fast constructor has not been built yet" from "built, and the # shape is not handled, so it is None and we fall back to validating". _UNSET = object() @@ -1095,7 +1191,86 @@ def _build_constructor( return namespace["_construct"] -class DataclassSchema[DataclassT](Schema): +class _FusedDataclassSchema(Schema): + """A ``Schema(All(inner_mapping, _Constructor))`` that compiles as one function. + + The shape ``create_dataclass_schema`` builds, for a top-level dataclass and for + every nested one. It validates interpreted through the fused engine + ``_install_fused_interpreted`` puts in place, and generates code for the mapping + and the construction together (``compile_mapping`` with ``construct``) when + compiled, whether by an explicit ``compile()``, the ``compile=True`` flag, or + the ``AUTO`` policy once it is hot. Compiling it compiles the nested dataclass + schemas under it first, so a tree generates all the way down. + """ + + # Set by ``_install_fused_interpreted``: the interpreted fused engine. + _fused_engine: CompiledSchema + + def __init__( + self, + schema: TypingAny, + *, + on_cycle: bool = False, + compile: bool | None = None, # noqa: A002 - mirrors ``Schema`` + ) -> None: + """Wrap the ``All``; ``on_cycle`` pins a schema on a recursion cycle interpreted.""" + # Before ``Schema.__init__``, which reads ``_compilable`` to decide arming. + self._on_cycle = on_cycle + super().__init__(schema, compile=compile) + + def _compilable(self) -> bool: + """Report that a dataclass compiles via its inner mapping, under the ``All``. + + A schema on a recursion cycle never does (see ``_mark_cycle``). + """ + return not self._on_cycle and isinstance( + self.schema.validators[0]._compiled, # noqa: SLF001 + _MappingValidator, + ) + + def _compile_from(self, interpreted: CompiledSchema) -> CompiledSchema: + """Compile the dataclass into one fused validate-and-construct function. + + The schema is ``All(inner_mapping, _Constructor)``. Compile the inner mapping + and splat the validated fields straight into the dataclass, so neither the + generic mapping loop nor the constructor's dict-comprehension runs. Any + failure bails to ``interpreted`` (the ``All``), which validates and + constructs the same way. An ``ALLOW_EXTRA`` dataclass is not generated + (``compile_mapping`` returns ``None``), since the constructor would reject + the unknown keys the mapping keeps. + + The nested dataclass schemas the builder made for the fields compile first. + Each is called from the generated code through its own ``Schema.__call__``, + so swapping its engine in place is all it takes; left alone they would stay + interpreted under every policy, since the generated outer calls them + directly and never through the bootstrap that promotes a hot schema. + """ + if self._on_cycle: + return interpreted + fused = self.__dict__.get("_fused_engine") + if fused is not None and interpreted is not fused: + # Already generated (a second ``compile()``): keep the one in place + # rather than stacking a new validator whose bail-out is the old one. + return interpreted + + inner, constructor = self.schema.validators + mapping = inner._compiled # noqa: SLF001 + if not isinstance( + mapping, _MappingValidator + ): # pragma: no cover - the inner schema is always a mapping + return interpreted + + for candidate in mapping._candidates: # noqa: SLF001 + for nested in _built_dataclass_schemas(candidate.value_schema): + nested.compile() + + generated = compile_mapping( + mapping, construct=constructor.dataclass_type, fallback=interpreted + ) + return generated if generated is not None else interpreted + + +class DataclassSchema[DataclassT](_FusedDataclassSchema): """A ``Schema`` generated from a dataclass type. Calling it validates a mapping against the dataclass's fields and returns a @@ -1123,6 +1298,7 @@ def __init__( additional_constraints, required=required, extra=extra, + _compile=compile, ) super().__init__(built.schema, compile=compile) # ``built`` carries its own fused engine, but this instance recompiled the @@ -1185,32 +1361,6 @@ def __call__( ) -> DataclassT: """Validate ``data`` and return the constructed dataclass instance.""" - def _compilable(self) -> bool: - """Report that a dataclass compiles via its inner mapping, under the ``All``.""" - return isinstance(self.schema.validators[0]._compiled, _MappingValidator) # noqa: SLF001 - - def _compile_from(self, interpreted: CompiledSchema) -> CompiledSchema: - """Compile the dataclass into one fused validate-and-construct function. - - The schema is ``All(inner_mapping, _Constructor)``. Compile the inner mapping - and splat the validated fields straight into the dataclass, so neither the - generic mapping loop nor the constructor's dict-comprehension runs. Any - failure bails to ``interpreted`` (the ``All``), which validates and - constructs the same way. An ``ALLOW_EXTRA`` dataclass is not generated - (``compile_mapping`` returns ``None``), since the constructor would reject - the unknown keys the mapping keeps. - """ - mapping = self.schema.validators[0]._compiled # noqa: SLF001 - if not isinstance( - mapping, _MappingValidator - ): # pragma: no cover - the inner schema is always a mapping - return interpreted - - generated = compile_mapping( - mapping, construct=self.dataclass_type, fallback=interpreted - ) - return generated if generated is not None else interpreted - def _typeddict_presence( annotation: TypingAny, *, default_required: bool @@ -1244,6 +1394,8 @@ def _typeddict_mapping( constraints: dict[str, TypingAny], self_refs: dict[type, _RecursiveSchemaRef], extra: int, + *, + compiled: bool | None, ) -> dict[TypingAny, TypingAny]: """Build the mapping schema for a TypedDict's fields.""" try: @@ -1269,7 +1421,9 @@ def _typeddict_mapping( required, field_type = _typeddict_presence( annotation, default_required=name in required_keys ) - value_schema = _annotation_to_schema(field_type, self_refs, extra) + value_schema = _annotation_to_schema( + field_type, self_refs, extra, compiled=compiled + ) if name in constraints: value_schema = All(value_schema, constraints[name]) @@ -1294,6 +1448,7 @@ def create_typeddict_schema( *, extra: int = PREVENT_EXTRA, _self_refs: dict[type, _RecursiveSchemaRef] | None = None, + _compile: bool | None = None, ) -> Schema: """Build a ``Schema`` that validates a mapping against a TypedDict. @@ -1318,10 +1473,14 @@ def create_typeddict_schema( ref = _RecursiveSchemaRef() self_refs[typeddict_type] = ref mapping = _typeddict_mapping( - typeddict_type, additional_constraints or {}, self_refs, extra + typeddict_type, + additional_constraints or {}, + self_refs, + extra, + compiled=_compile, ) - inner = Schema(mapping, extra=extra) + inner = Schema(mapping, extra=extra, compile=_compile) ref.bind(inner) return inner @@ -1355,6 +1514,7 @@ def __init__( typeddict_type, additional_constraints, extra=extra, + _compile=compile, ) # ``built.schema`` is the raw mapping, so ``extra`` has to be passed again # when re-wrapping; the dataclass path gets away without it because there diff --git a/src/probatio/decorator.py b/src/probatio/decorator.py index 1db8fccf..d52694ce 100644 --- a/src/probatio/decorator.py +++ b/src/probatio/decorator.py @@ -148,7 +148,9 @@ def _parameter_schema( if parameter.kind not in _VALIDATED_KINDS: continue if name in hints: - inferred = _annotation_to_schema(hints[name], {}, PREVENT_EXTRA) + inferred = _annotation_to_schema( + hints[name], {}, PREVENT_EXTRA, compiled=None + ) _reject_validator_annotation( func, f"parameter {name!r}", hints[name], inferred ) @@ -182,7 +184,9 @@ def _return_schema( f"probatio: returns=True needs a return annotation on {_name(func)!r}" ) raise SchemaError(message) - inferred = _annotation_to_schema(hints[_RETURN], {}, PREVENT_EXTRA) + inferred = _annotation_to_schema( + hints[_RETURN], {}, PREVENT_EXTRA, compiled=None + ) _reject_validator_annotation(func, "the return", hints[_RETURN], inferred) return Schema(inferred) return Schema(returns) diff --git a/tests/__snapshots__/test_public_surface.ambr b/tests/__snapshots__/test_public_surface.ambr index c9517b92..fa323fac 100644 --- a/tests/__snapshots__/test_public_surface.ambr +++ b/tests/__snapshots__/test_public_surface.ambr @@ -1834,6 +1834,7 @@ 'required KEYWORD_ONLY False', 'extra KEYWORD_ONLY 0', '_self_refs KEYWORD_ONLY None', + '_compile KEYWORD_ONLY None', ]), }), 'create_typeddict_schema': dict({ @@ -1843,6 +1844,7 @@ 'additional_constraints POSITIONAL_OR_KEYWORD None', 'extra KEYWORD_ONLY 0', '_self_refs KEYWORD_ONLY None', + '_compile KEYWORD_ONLY None', ]), }), 'current_context': dict({ diff --git a/tests/test_codegen.py b/tests/test_codegen.py index a7cb5453..54469f69 100644 --- a/tests/test_codegen.py +++ b/tests/test_codegen.py @@ -8,10 +8,11 @@ from __future__ import annotations +import time from dataclasses import dataclass, field from enum import Enum, StrEnum from types import MappingProxyType -from typing import TypedDict +from typing import Annotated, TypedDict import pytest @@ -43,6 +44,7 @@ from probatio._codegen import compile_mapping from probatio._engine import _MappingValidator from probatio.codecs.jsonschema import _DeferredRef +from probatio.dataclass_schema import _built_dataclass_schemas from probatio.schema import _AUTO_COMPILE_THRESHOLD @@ -747,3 +749,198 @@ def reentrant(self: Schema, interpreted: object) -> object: assert schema(data) == data assert arrived, "the re-entrant call never happened, so nothing was proven" + + +# --- nested dataclass schemas compile with the schema that holds them --- + + +@dataclass +class _Unit: + """A unit/value pair, the leaf of a dataclass tree.""" + + unit: str + value: float + + +def _pass_through(unit: _Unit) -> _Unit: + """A plain callable constraint: no child-schema protocol, nothing to walk.""" + return unit + + +_OWN_UNIT = DataclassSchema(_Unit, compile=False) + + +@dataclass +class _Reading: + """A dataclass tree holding ``_Unit`` through every shape the compile walk enters. + + ``own`` carries a ``DataclassSchema`` the user built with its own flag; the walk + must leave that one alone. + """ + + direct: _Unit + maybe: _Unit | None + items: list[_Unit] + by_name: dict[str, _Unit] + constrained: Annotated[_Unit, _pass_through] + own: Annotated[_Unit, _OWN_UNIT] + + +_UNIT = {"unit": "W", "value": 1.5} +_READING = { + "direct": _UNIT, + "maybe": None, + "items": [_UNIT, _UNIT], + "by_name": {"a": _UNIT}, + "constrained": _UNIT, + "own": _UNIT, +} + + +def _nested_schemas(schema: DataclassSchema[object]) -> list[Schema]: + """The builder-made dataclass schemas under a ``DataclassSchema``'s fields.""" + mapping = schema.schema.validators[0].schema + return [ + nested + for value in mapping.values() + for nested in _built_dataclass_schemas(value) + ] + + +@pytest.mark.parametrize( + "data", + [ + _READING, + {**_READING, "maybe": _UNIT}, + {**_READING, "direct": {"unit": "W", "value": "x"}}, + {**_READING, "items": [_UNIT, {"unit": 1, "value": 1.0}]}, + {**_READING, "by_name": {"a": {"unit": "W"}}}, + {**_READING, "own": {"unit": "W", "value": 1.5, "extra": 1}}, + ], +) +def test_nested_dataclass_schemas_compile_with_the_outer( + data: dict[str, object], +) -> None: + """Compiling a dataclass compiles the nested dataclass schemas under its fields. + + Six nested schemas sit under ``_Reading``: a direct field, a ``Maybe``, a list, + a dict, an ``Annotated`` with a constraint, and the base of the ``own`` field. + Every one generates; the ``DataclassSchema`` the user placed in ``own`` keeps + its ``compile=False``. Results and errors match the interpreted tree. + """ + interpreted = DataclassSchema(_Reading, compile=False) + compiled = DataclassSchema(_Reading, compile=True).compile() + assert _is_compiled(compiled) + nested = _nested_schemas(compiled) + assert len(nested) == 6 + assert all(_is_compiled(schema) for schema in nested) + assert not _is_compiled(_OWN_UNIT) + assert _outcome(interpreted, data) == _outcome(compiled, data) + + +def test_compile_false_keeps_the_nested_dataclass_schemas_interpreted() -> None: + """``compile=False`` on the outer dataclass pins the whole tree interpreted. + + The flag is threaded into every nested schema the builder makes, so under the + ON policy none of them arms or generates; the interpreted benchmark baseline + for a dataclass tree stays an interpreted baseline. + """ + set_compile_policy(CompilePolicy.ON) + try: + schema = DataclassSchema(_Reading, compile=False) + for _ in range(3): + assert schema(_READING).direct == _Unit("W", 1.5) + assert not _is_compiled(schema) + for nested in _nested_schemas(schema): + assert not _is_compiled(nested) + assert "_interpreted" not in nested.__dict__ + finally: + set_compile_policy(CompilePolicy.OFF) + + +def test_dataclass_compile_is_idempotent() -> None: + """A second ``compile()`` keeps the generated tree instead of stacking a new one.""" + schema = DataclassSchema(_Reading).compile() + before = [schema._compiled, *(n._compiled for n in _nested_schemas(schema))] + schema.compile() + after = [schema._compiled, *(n._compiled for n in _nested_schemas(schema))] + assert after == before + + +def test_auto_compiles_a_hot_dataclass_tree() -> None: + """Under AUTO a dataclass that goes hot compiles its nested schemas with it.""" + set_compile_policy(CompilePolicy.AUTO) + try: + schema = DataclassSchema(_Reading) + for _ in range(_AUTO_COMPILE_THRESHOLD): + assert schema(_READING).items[1] == _Unit("W", 1.5) + assert _is_compiled(schema) + assert all(_is_compiled(nested) for nested in _nested_schemas(schema)) + finally: + set_compile_policy(CompilePolicy.OFF) + + +@dataclass +class _Ping: + """One half of a mutually recursive pair.""" + + pong: _Pong | None = None + + +@dataclass +class _Pong: + """The other half, closing the cycle back to ``_Ping``.""" + + ping: _Ping | None = None + + +def test_recursion_cycle_stays_interpreted_while_the_top_compiles() -> None: + """The schemas on a recursion cycle never compile; the schema above it does. + + ``_Ping``'s nested ``_Pong`` refers back to ``_Ping``, so both are on the cycle + and are pinned interpreted even when the outer ``compile()`` walks into them. + The top-level ``DataclassSchema`` sits above the cycle and still generates. + """ + interpreted = DataclassSchema(_Ping, compile=False) + compiled = DataclassSchema(_Ping, compile=True).compile() + assert _is_compiled(compiled) + (pong,) = _nested_schemas(compiled) + assert pong._on_cycle + assert not _is_compiled(pong) + + data: dict[str, object] = {"pong": {"ping": {"pong": {"ping": None}}}} + assert _outcome(interpreted, data) == _outcome(compiled, data) + bad: dict[str, object] = {"pong": {"ping": {"pong": {"ping": 1}}}} + assert _outcome(interpreted, bad) == _outcome(compiled, bad) + + +@dataclass +class _Chain: + """A self-referential dataclass, for the deep-failure regression test.""" + + value: int + nxt: _Chain | None = None + + +def test_failing_deep_recursion_costs_linear_time() -> None: + """A failure deep in a recursive tree is validated once per level, not 2**depth. + + A generated validator bails to its interpreted twin on a failure, and the twin + re-enters the generated code below it. If the schemas on a recursion cycle + compiled, each level of data would double the work, and a failure 24 levels + down would take minutes; pinning the cycle interpreted keeps it linear. Under + the ON policy with everything that may compile compiled. + """ + set_compile_policy(CompilePolicy.ON) + try: + depth = 24 + data: dict[str, object] = {"value": "bottom"} + for index in range(depth): + data = {"value": index, "nxt": data} + schema = DataclassSchema(_Chain) + start = time.perf_counter() + with pytest.raises(MultipleInvalid, match="expected int"): + schema(data) + assert time.perf_counter() - start < 1.0 + finally: + set_compile_policy(CompilePolicy.OFF) From 8c09928ee8662ae8dedceb08f9ea16b0ff468c1f Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 19:33:22 +0000 Subject: [PATCH 2/2] Keep a recursive TypedDict's cycle interpreted too The nested schema for a TypedDict on a recursion cycle could enter the generated path: under the ON policy already, and now also through the compile=True flag the outer schema threads down. That is the same cascade the dataclass rule guards against, a failing validation exponential in the data's depth (about 13 seconds at 20 levels). The inner schema is now built with compile=False when its reference closes a cycle, so it never arms under any policy; the top-level TypedDictSchema above the cycle still compiles. A regression test pins a failure 24 levels deep to linear time. Raised by the CodeRabbit review on the pull request. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ANcsDvtNPxCRTk33uagPak --- .../content/docs/guides/compiled-schemas.md | 4 +- src/probatio/dataclass_schema.py | 7 +++- tests/test_codegen.py | 38 ++++++++++++++++++- 3 files changed, 45 insertions(+), 4 deletions(-) diff --git a/docs/src/content/docs/guides/compiled-schemas.md b/docs/src/content/docs/guides/compiled-schemas.md index 97f13091..5c1de2d4 100644 --- a/docs/src/content/docs/guides/compiled-schemas.md +++ b/docs/src/content/docs/guides/compiled-schemas.md @@ -127,8 +127,8 @@ inlines the validators that show up most: `Coerce` to `int`/`float`, numeric validated value), and `All`/`Any` chains of those. Anything it does not handle is simply left interpreted. A recursive schema (`Self`, -a recursive JSON Schema `$ref`, or a dataclass on a recursion cycle: one that holds -itself, or two that hold each other), an exclusive or inclusive group, a multi-element +a recursive JSON Schema `$ref`, or a dataclass or TypedDict on a recursion cycle: +one that holds itself, or two that hold each other), an exclusive or inclusive group, a multi-element list (an "each item matches one of these" union), a list whose element does not inline, an exotic key: the schema works exactly as before, just without the speedup. You never lose correctness by compiling; at worst you do not gain speed. diff --git a/src/probatio/dataclass_schema.py b/src/probatio/dataclass_schema.py index 403f5384..cbe5bb53 100644 --- a/src/probatio/dataclass_schema.py +++ b/src/probatio/dataclass_schema.py @@ -1480,7 +1480,12 @@ def create_typeddict_schema( compiled=_compile, ) - inner = Schema(mapping, extra=extra, compile=_compile) + # A TypedDict on a recursion cycle is pinned interpreted, for the reason a + # dataclass on one is (see ``_mark_cycle``): a generated validator bails to its + # interpreted twin on a failure, which re-enters the generated code below it, + # so a compiled schema on a cycle makes a failing validation exponential in + # the data's depth. ``compile=False`` never arms, under any policy. + inner = Schema(mapping, extra=extra, compile=False if ref.on_cycle else _compile) ref.bind(inner) return inner diff --git a/tests/test_codegen.py b/tests/test_codegen.py index 54469f69..9ff0b34f 100644 --- a/tests/test_codegen.py +++ b/tests/test_codegen.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field from enum import Enum, StrEnum from types import MappingProxyType -from typing import Annotated, TypedDict +from typing import Annotated, NotRequired, TypedDict import pytest @@ -944,3 +944,39 @@ def test_failing_deep_recursion_costs_linear_time() -> None: assert time.perf_counter() - start < 1.0 finally: set_compile_policy(CompilePolicy.OFF) + + +class _Link(TypedDict): + """A self-referential TypedDict, for the deep-failure regression test.""" + + value: int + nxt: NotRequired[_Link | None] + + +def test_recursive_typeddict_keeps_the_cycle_interpreted() -> None: + """A TypedDict on a recursion cycle never compiles; the schema above it does. + + The inner schema the builder makes for ``_Link`` refers back to itself, so it + is pinned interpreted under the ``compile=True`` flag and the ON policy alike. + The top-level ``TypedDictSchema`` sits above the cycle and still generates, so + a failure deep in the data is validated once per level, not 2**depth. + """ + set_compile_policy(CompilePolicy.ON) + try: + schema = TypedDictSchema(_Link, compile=True).compile() + assert _is_compiled(schema) + # The recursive reference holds the inner schema's engine: the interpreted + # mapping validator when pinned, a generated function if it had compiled. + ref = schema.schema["nxt"].validator + assert isinstance(ref._validate, _MappingValidator) + + depth = 24 + data: dict[str, object] = {"value": "bottom"} + for index in range(depth): + data = {"value": index, "nxt": data} + start = time.perf_counter() + with pytest.raises(MultipleInvalid, match="expected int"): + schema(data) + assert time.perf_counter() - start < 1.0 + finally: + set_compile_policy(CompilePolicy.OFF)