From 42d831e501ccece73d017399ec73bfa0ed931087 Mon Sep 17 00:00:00 2001 From: Alessandro Cecchini Date: Tue, 8 Sep 2026 11:42:17 +0200 Subject: [PATCH 1/2] docs: plan useful static Tree container support --- plans/2026-09-08-static-tree-containers.md | 87 ++++++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 plans/2026-09-08-static-tree-containers.md diff --git a/plans/2026-09-08-static-tree-containers.md b/plans/2026-09-08-static-tree-containers.md new file mode 100644 index 0000000..a86f338 --- /dev/null +++ b/plans/2026-09-08-static-tree-containers.md @@ -0,0 +1,87 @@ +# Give ordinary Tree consumers a useful static type + + +Maintain this ExecPlan according to `PLANS.md`. This independent PR addresses the Tree portion of audit A06 and is stacked on checker harness PR #17. Runtime tree traversal and structure binding remain unchanged. + +## Purpose / Big Picture + + +Tree[int] should accept integer leaves and existing typed lists, tuples and dictionaries of integer leaves. It should reject strings and wrong nested leaves. The present nominal stub rejects real containers. Replace it only after demonstrating a model that retains leaf information in all four supported checkers. + +## Progress + + +- [x] (2026-09-08) Created feature worktree and tested recursive-container/protocol prototypes with all four checkers. +- [ ] Open draft PR and reproduce the nominal stub's real-call failures. +- [ ] Implement the smallest proven shared static model. +- [ ] Verify nested valid/invalid containers, inference, strings and custom-node boundaries. +- [ ] Run all checker/Python/floor and focused runtime checks. +- [ ] Update docs, changelog and evidence with explicit support limits. + +## Surprises & Discoveries + + +A recursive alias using list/dict directly rejects already-typed containers because their element types are invariant. In the same prototype ty accepts even deliberate errors. A sequence/protocol model can admit strings through recursive iteration: type stubs expose inherited sequence behavior that is not equivalent to Python's runtime attributes. Another protocol version rejects direct strings but mypy accepts a list of strings. + +The current successful small probe uses private covariant protocols for list, tuple and mapping behavior, with a nonrecursive outer union. List's pop result carries recursive leaf information; tuple iteration and tuple concatenation distinguish it from self-iterating strings; mapping values carry recursive leaf information without constraining keys. All four engines accept six valid pretyped cases and reject both deliberate wrong cases. These protocol members are descriptive only; validation does not call them or mutate inputs. The full implementation must extend the probe to mixed nesting, NumPy leaves and additional invalid forms before promotion. + +## Decision Log + + +Decision: Preserve Tree[Leaf] syntax and share the model in the TYPE_CHECKING section of `src/bearshape/_tree.py`, re-exporting it from optree/JAX. Rationale: the two public tree backends should not carry divergent fake nominal classes or duplicate static definitions. No runtime dependency or new module is needed. Date: 2026-09-08. + +Decision: Establish tested static support for ordinary lists, tuples, dictionaries, leaves and None; keep arbitrary backend registration a runtime property. Rationale: Python static typing cannot infer a dynamically modified pytree registry. For custom nodes, document the existing TYPE_CHECKING alias pattern using the user's concrete node type and the runtime Tree annotation. This preserves existing runtime support without introducing a new public form or claiming every structural match is registered. + +## Outcomes & Retrospective + + +Implementation pending. Do not adopt a model merely because valid examples pass: exact negative diagnostics, especially direct/nested strings and invalid array dtypes, are required. State backend registration limits accurately. Structure-bearing Tree syntax remains runtime-only. + +## Context and Orientation + + +Worktree `/Users/ale/Code/bearshape-worktrees/static-tree-containers`, branch `codex/static-tree-containers`, base `3eac024`. `src/bearshape/_tree.py` owns runtime traversal and structure checking. The TYPE_CHECKING branches in `optree.py` and `jax.py` each define a nominal Tree class. `tests/typing/` and `tests/typing_negative/` use the four-engine harness from PR #17. `tests/test_tree.py` owns runtime tree regressions; add an executable consumer fixture and a maintained runtime invocation there. + +## Plan of Work + + +Add real positive consumer calls with typed leaves, list[int], list[list[int]], tuple mixtures, dict[str, list[int]], None, empty containers, and nested NumPy arrays. Add negative wrong scalars, nested strings, sets and incompatible NumPy dtype examples for both optree/JAX aliases. Save the failing-before checker results. + +Implement private covariant protocols and the shared static Tree alias under TYPE_CHECKING in `_tree.py`; replace the public nominal classes with imports of that alias. Use the smallest protocol member sets demonstrated by the prototypes. Extend inference checks to ensure a Tree annotation does not collapse to Any/Unknown. Execute positive consumer functions decorated with beartype so the same examples establish runtime acceptance. Keep custom-node registration and structure binding unchanged and covered by existing tests. + +Document the ordinary-container model, the runtime registry boundary and the existing conditional-alias technique for custom registered node types. Do not imply that a structurally matching arbitrary class is registered automatically. Update CHANGELOG. Repeat the complete checker contract on Python 3.10–3.14 and maintained floor engines; run runtime tree tests with exact rc0 endpoints, hooks and locked dev coverage. + +## Concrete Steps + + +From this worktree: + + uv sync --locked + uv run --locked pytest tests/test_typecheck.py -n 4 + uv run --locked pytest tests/test_tree.py -n 4 + uv run --locked tox run -e dev + uv run --locked prek run -a + +Before implementation preserve the nominal-stub failures. Use interpreter-matched environments for the other Python versions and the existing exact-rc0 CPU environments for runtime source checks. The integrated artifact milestone must repeat consumer checking against the installed wheel. + +## Validation and Acceptance + + +Every checker accepts the maintained valid pretyped containers and rejects all marked errors at the intended source line/category. None and empty containers match actual backend flattening. NumPy leaf annotations retain useful dtype information. Runtime-decorated consumers and existing tree tests pass; no runtime mutation, registry change, Any-valued leaf, or broad new checker suppression is introduced. Custom nodes retain their existing runtime path and have an explicit static spelling based on the user's concrete node type. + +## Idempotence and Recovery + + +Keep prototypes outside the source package and preserve their counterexamples. Only promote the model after complete negative/inference proof. Work remains isolated from other PRs until integration. Do not merge without user validation. + +## Artifacts and Notes + + +Prototypes are under `/Users/ale/Code/bearshape-implementation-2026-09-08/evidence/tree-typing-prototypes/`: `recursive_containers.py`, `recursive_protocols.py`, `reverse_protocols.py`, and `pop_protocols.py`, with per-engine outputs. The first three demonstrate why acceptance-only validation is inadequate. Save production before/after and matrix evidence as `tree-typing-*.log` in the evidence directory. + +## Interfaces and Dependencies + + +Keep public `bearshape.optree.Tree` and `bearshape.jax.Tree` subscriptions unchanged. The static alias describes leaves and supported container behavior; runtime still uses _TreeFactory and the backend registry. Use standard typing/collections protocols and the existing checker harness. No new runtime or development dependency is required. + +Revision note — 2026-09-08: Recorded prototype counterexamples and focused Tree implementation plan before source changes. From 2513af46cb986c1936fa1f87a7159725e8ffde2c Mon Sep 17 00:00:00 2001 From: Alessandro Cecchini Date: Tue, 8 Sep 2026 12:00:15 +0200 Subject: [PATCH 2/2] fix: preserve leaf types in ordinary static Tree containers --- CHANGELOG.md | 5 + docs/features/static-typing.md | 37 ++++-- docs/features/tree-annotations.md | 15 +++ plans/2026-09-08-static-tree-containers.md | 20 +-- src/bearshape/_tree.py | 35 +++++- src/bearshape/jax.py | 7 +- src/bearshape/optree.py | 7 +- tests/test_tree.py | 25 ++++ tests/typing/check_tree.py | 2 +- tests/typing/check_tree_consumers.py | 128 ++++++++++++++++++++ tests/typing_negative/invalid_tree_calls.py | 52 ++++++++ 11 files changed, 298 insertions(+), 35 deletions(-) create mode 100644 tests/typing/check_tree_consumers.py create mode 100644 tests/typing_negative/invalid_tree_calls.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ceaae1..b798ce9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,9 +8,14 @@ and this project follows ## [Unreleased] +- Replace nominal static Tree stubs with a shared model for ordinary typed + leaves, lists, tuples and dictionaries, preserving expected leaf errors. + Document the concrete-type alias pattern for custom JAX nodes. + - Verify real consumer calls, inferred types, and expected errors with pyright, mypy, ty, and pyrefly; require selected tools instead of silently skipping them. + - Update the locked checker versions and make `check(conf=None)` match its public overloads. diff --git a/docs/features/static-typing.md b/docs/features/static-typing.md index 114bed1..166c48e 100644 --- a/docs/features/static-typing.md +++ b/docs/features/static-typing.md @@ -181,24 +181,39 @@ keep in their toolbox first. ## Tree annotations -`Tree` has a split contract: +`Tree[Leaf]` supports ordinary leaves, lists, tuples (including named tuples), +dictionaries and None. Existing typed containers such as `list[int]` and +`dict[str, list[int]]` can be passed to `Tree[int]`. Wrong leaves, including +strings hidden inside a numeric tree, are checked by the consumer fixtures. An +empty container or None has no leaves for the default backend traversal. -- `Tree[F32[N, C]]` is checker-friendly and tested -- `Tree[F32[N], T]`, `Tree[F32[N], T, ...]`, and similar structure-bearing forms - are runtime-only and need a targeted ignore +The static model describes container behavior; it cannot infer backend node +registration or tree structure. The selected backend must actually recognize a +custom container. For a registered JAX node, keep its concrete static type with +the existing conditional-alias pattern: ```python -from bearshape import N, T -from bearshape.numpy import F32 -from bearshape.optree import Tree +from typing import TYPE_CHECKING, TypeAlias -def leaves_only(x: Tree[F32[N]]) -> Tree[F32[N]]: - return x +from bearshape import N +from bearshape.jax import Tree +from bearshape.numpy import F32 -def structure_checked(x: Tree[F32[N], T]) -> Tree[F32[N]]: # type: ignore[valid-type] - return x +# Batch is your concrete class, registered with jax.tree_util. +if TYPE_CHECKING: + BatchTree: TypeAlias = Batch +else: + BatchTree = Tree[F32[N]] ``` +The executable fixture `tests/typing/check_tree_consumers.py` contains the +complete registered class and a decorated consumer. Tree structure arguments +such as `Tree[F32[N], T]` remain runtime-only; use a checker-only leaf alias as +shown above when you need a named structure constraint. + +The optree backend uses its default registry. A class registered only in an +optree namespace is not automatically recognized by this Tree annotation. + ## Backend notes The typing model differs slightly from runtime behavior: diff --git a/docs/features/tree-annotations.md b/docs/features/tree-annotations.md index 6170512..0999c78 100644 --- a/docs/features/tree-annotations.md +++ b/docs/features/tree-annotations.md @@ -168,3 +168,18 @@ only | | `Tree[LeafType, T]` | Full structure binding | | Bottom-level only | | `Tree[LeafType, T, S]` | T = top (one level), S = full remaining | | `Tree[LeafType, T, S, ...]` | T = top, S = next, inner unchecked | | `Tree[LeafType, ..., T, S]` | S = bottom, T = second-from-bottom | + +## Static container support and registration + +`Tree[Leaf]` models ordinary leaves, lists, tuples (including named tuples), +dictionaries and None. Typed variables such as `list[int]` and +`dict[str, list[int]]` are valid inputs to `Tree[int]`; wrong nested leaves +remain checker errors. Validation uses the backend registry and does not mutate +the input containers. + +Static structural compatibility cannot establish that a custom class is +registered. For custom JAX nodes, use the concrete node type in a TYPE_CHECKING +alias and the Tree annotation at runtime; the +[static typing guide](static-typing.md#tree-annotations) describes this existing +pattern. The optree backend uses the default registry, so a class registered +only in an optree namespace is not automatically traversed by this annotation. diff --git a/plans/2026-09-08-static-tree-containers.md b/plans/2026-09-08-static-tree-containers.md index a86f338..614621d 100644 --- a/plans/2026-09-08-static-tree-containers.md +++ b/plans/2026-09-08-static-tree-containers.md @@ -12,30 +12,30 @@ Tree[int] should accept integer leaves and existing typed lists, tuples and dict - [x] (2026-09-08) Created feature worktree and tested recursive-container/protocol prototypes with all four checkers. -- [ ] Open draft PR and reproduce the nominal stub's real-call failures. -- [ ] Implement the smallest proven shared static model. -- [ ] Verify nested valid/invalid containers, inference, strings and custom-node boundaries. -- [ ] Run all checker/Python/floor and focused runtime checks. -- [ ] Update docs, changelog and evidence with explicit support limits. +- [x] (2026-09-08) Opened PR #20 and reproduced real-call failures with the old nominal stubs. +- [x] (2026-09-08) Replaced both nominal stubs with one shared static model under TYPE_CHECKING. +- [x] (2026-09-08) Verified ordinary pretyped containers, named tuples, arrays, empty inputs, strings, fourteen negative sites per checker, and the concrete custom JAX node alias. +- [x] (2026-09-08) All four engines pass on Python 3.10–3.14 and floor lanes. Exact rc0 endpoints each pass 108 tree tests. Locked dev tox: 1,044 passed, five expected skips, 91.24% coverage. +- [x] (2026-09-08) Updated docs/changelog and recorded default-registry limits. Hooks pass. ## Surprises & Discoveries A recursive alias using list/dict directly rejects already-typed containers because their element types are invariant. In the same prototype ty accepts even deliberate errors. A sequence/protocol model can admit strings through recursive iteration: type stubs expose inherited sequence behavior that is not equivalent to Python's runtime attributes. Another protocol version rejects direct strings but mypy accepts a list of strings. -The current successful small probe uses private covariant protocols for list, tuple and mapping behavior, with a nonrecursive outer union. List's pop result carries recursive leaf information; tuple iteration and tuple concatenation distinguish it from self-iterating strings; mapping values carry recursive leaf information without constraining keys. All four engines accept six valid pretyped cases and reject both deliberate wrong cases. These protocol members are descriptive only; validation does not call them or mutate inputs. The full implementation must extend the probe to mixed nesting, NumPy leaves and additional invalid forms before promotion. +The current successful small probe uses private covariant protocols for list, tuple and mapping behavior, with a nonrecursive outer union. List's pop result carries recursive leaf information; tuple indexing and tuple concatenation distinguish it from self-iterating strings; mapping values carry recursive leaf information without constraining keys. All four engines accept six valid pretyped cases and reject both deliberate wrong cases. These protocol members are descriptive only; validation does not call them or mutate inputs. The full implementation must extend the probe to mixed nesting, NumPy leaves and additional invalid forms before promotion. ## Decision Log -Decision: Preserve Tree[Leaf] syntax and share the model in the TYPE_CHECKING section of `src/bearshape/_tree.py`, re-exporting it from optree/JAX. Rationale: the two public tree backends should not carry divergent fake nominal classes or duplicate static definitions. No runtime dependency or new module is needed. Date: 2026-09-08. +Decision: Preserve Tree[Leaf] syntax and share the model in the TYPE_CHECKING section of `src/bearshape/_tree.py`, re-exporting it from optree/JAX. Rationale: the two public tree backends should not carry divergent fake nominal classes or duplicate static definitions. No runtime dependency or new module is needed. The existing typing_extensions.TypeAliasType represents the outer alias so ty also retains existing checker-only aliases built from Tree. Date: 2026-09-08. Decision: Establish tested static support for ordinary lists, tuples, dictionaries, leaves and None; keep arbitrary backend registration a runtime property. Rationale: Python static typing cannot infer a dynamically modified pytree registry. For custom nodes, document the existing TYPE_CHECKING alias pattern using the user's concrete node type and the runtime Tree annotation. This preserves existing runtime support without introducing a new public form or claiming every structural match is registered. ## Outcomes & Retrospective -Implementation pending. Do not adopt a model merely because valid examples pass: exact negative diagnostics, especially direct/nested strings and invalid array dtypes, are required. State backend registration limits accurately. Structure-bearing Tree syntax remains runtime-only. +Implemented real ordinary-container acceptance and maintained exact negative diagnostics for direct/nested strings and invalid array dtypes. The same positive consumer executes with beartype, including a custom registered JAX node; a separate runtime test rejects that node with the wrong leaf dtype. State backend registration limits accurately. Structure-bearing Tree syntax remains runtime-only. ## Context and Orientation @@ -77,7 +77,7 @@ Keep prototypes outside the source package and preserve their counterexamples. O ## Artifacts and Notes -Prototypes are under `/Users/ale/Code/bearshape-implementation-2026-09-08/evidence/tree-typing-prototypes/`: `recursive_containers.py`, `recursive_protocols.py`, `reverse_protocols.py`, and `pop_protocols.py`, with per-engine outputs. The first three demonstrate why acceptance-only validation is inadequate. Save production before/after and matrix evidence as `tree-typing-*.log` in the evidence directory. +Prototypes are under `/Users/ale/Code/bearshape-implementation-2026-09-08/evidence/tree-typing-prototypes/`: `recursive_containers.py`, `recursive_protocols.py`, `reverse_protocols.py`, and `pop_protocols.py`, with per-engine outputs. The first three demonstrate why acceptance-only validation is inadequate. Production evidence uses `tree-typing-*.log`: `before` captures the nominal-stub failures; `focused` captures the initial 115 passing tree/checker tests; `tox` captures 1,044 passing dev tests, 91.24% coverage and all floor engines; `python-3.11` through `python-3.14` capture the interpreter-matched checker matrix; `rc0-py310` and `rc0-py314` each report 108 passing tree tests; `hooks-final` is clean. The final model uses tuple indexing because pinned pyrefly describes named-tuple iteration as Iterable rather than Iterator. A namespaced optree custom registration is not visible to the current default-registry Tree; the maintained custom-node example is explicitly JAX. ## Interfaces and Dependencies @@ -85,3 +85,5 @@ Prototypes are under `/Users/ale/Code/bearshape-implementation-2026-09-08/eviden Keep public `bearshape.optree.Tree` and `bearshape.jax.Tree` subscriptions unchanged. The static alias describes leaves and supported container behavior; runtime still uses _TreeFactory and the backend registry. Use standard typing/collections protocols and the existing checker harness. No new runtime or development dependency is required. Revision note — 2026-09-08: Recorded prototype counterexamples and focused Tree implementation plan before source changes. + +Revision note — 2026-09-08: Implemented and validated the shared static model, corrected the prototype for named tuples and existing aliases, and documented the actual registry boundary. diff --git a/src/bearshape/_tree.py b/src/bearshape/_tree.py index 049eb1f..029bba5 100644 --- a/src/bearshape/_tree.py +++ b/src/bearshape/_tree.py @@ -16,8 +16,9 @@ Structure arguments (``T``, ``S``, ``...``) are **runtime-only**. Type checkers see ``Tree`` as ``Tree[LeafType]`` (one type parameter) and cannot validate multi-arg structure syntax like ``Tree[F32[N], T]``. - Leaf-only annotations such as ``Tree[F32[N, C]]`` are fully supported - by all type checkers. + Leaf-only annotations such as ``Tree[F32[N, C]]`` model ordinary + leaves, lists, tuples and dictionaries. Custom node registration is a + runtime property; use the concrete node type in a checker-only alias. Import ``Tree`` from an explicit backend module:: @@ -67,6 +68,36 @@ def f(x: Tree[int, T], y: Tree[int, S], z: Tree[int, T, S]): ... make_runtime_hint, ) +if tp.TYPE_CHECKING: + from collections.abc import ValuesView + + from typing_extensions import TypeAliasType + + _Leaf_co = tp.TypeVar("_Leaf_co", covariant=True) + + # These members describe containers; validation never invokes them. In + # particular, self-iterating strings must not satisfy the recursive model. + class _TreeList(tp.Protocol[_Leaf_co]): + def pop(self, index: int = -1, /) -> _StaticTree[_Leaf_co]: ... + + class _TreeTuple(tp.Protocol[_Leaf_co]): + def __getitem__(self, index: int, /) -> _StaticTree[_Leaf_co]: ... + def __add__(self, value: tuple[object, ...], /) -> tuple[object, ...]: ... + + class _TreeMapping(tp.Protocol[_Leaf_co]): + def values(self) -> ValuesView[_StaticTree[_Leaf_co]]: ... + + _StaticTree = TypeAliasType( + "_StaticTree", + _Leaf_co + | _TreeList[_Leaf_co] + | _TreeTuple[_Leaf_co] + | _TreeMapping[_Leaf_co] + | None, + type_params=(_Leaf_co,), + ) + + __all__ = ["S", "Structure", "T"] diff --git a/src/bearshape/jax.py b/src/bearshape/jax.py index 8ac0cb1..4d2c5f3 100644 --- a/src/bearshape/jax.py +++ b/src/bearshape/jax.py @@ -374,12 +374,7 @@ def _get_jax_tree_util() -> tp.Any: if tp.TYPE_CHECKING: - _T = tp.TypeVar("_T") - - class Tree(tp.Generic[_T]): - """Static type stub — ``Tree[LeafType]`` for type checkers.""" - - def __class_getitem__(cls, item: object) -> type: ... + from ._tree import _StaticTree as Tree else: Tree = _TreeFactory(_get_jax_tree_util, name="Tree") diff --git a/src/bearshape/optree.py b/src/bearshape/optree.py index d65710f..1bc7388 100644 --- a/src/bearshape/optree.py +++ b/src/bearshape/optree.py @@ -36,12 +36,7 @@ def _get_optree() -> tp.Any: if tp.TYPE_CHECKING: - _T = tp.TypeVar("_T") - - class Tree(tp.Generic[_T]): - """Static type stub — ``Tree[LeafType]`` for type checkers.""" - - def __class_getitem__(cls, item: object) -> type: ... + from ._tree import _StaticTree as Tree else: Tree = _TreeFactory(_get_optree, name="Tree") diff --git a/tests/test_tree.py b/tests/test_tree.py index 83d7359..64f3fc5 100644 --- a/tests/test_tree.py +++ b/tests/test_tree.py @@ -1473,3 +1473,28 @@ def test_optree_tree_repr(self) -> None: from bearshape.optree import Tree as OptreeTree assert repr(OptreeTree) == "Tree" + + +def test_typed_tree_consumer_examples() -> None: + """Run the same real container/custom-node calls checked by every engine.""" + import subprocess + import sys + from pathlib import Path + + pytest.importorskip("jax") + subprocess.run( + [sys.executable, "-m", "tests.typing.check_tree_consumers"], + check=True, + cwd=Path(__file__).resolve().parents[1], + capture_output=True, + text=True, + timeout=60, + ) + + +def test_custom_jax_node_leaf_dtype_is_checked() -> None: + pytest.importorskip("jax") + from tests.typing.check_tree_consumers import Batch, custom_node + + with pytest.raises(BeartypeCallHintParamViolation): + custom_node(Batch(np.ones(3, dtype=np.int32))) diff --git a/tests/typing/check_tree.py b/tests/typing/check_tree.py index c0270eb..7648903 100644 --- a/tests/typing/check_tree.py +++ b/tests/typing/check_tree.py @@ -15,7 +15,7 @@ # Import validation # --------------------------------------------------------------------------- -_pt = Tree # Tree should be importable +_pt = Tree[int] # Tree should be importable and subscriptable _st = Structure # Structure should be importable # --------------------------------------------------------------------------- diff --git a/tests/typing/check_tree_consumers.py b/tests/typing/check_tree_consumers.py new file mode 100644 index 0000000..63f0c87 --- /dev/null +++ b/tests/typing/check_tree_consumers.py @@ -0,0 +1,128 @@ +"""Real container calls shared by the optree and JAX tree annotations.""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import TYPE_CHECKING, NamedTuple, TypeAlias + +import numpy as np +from beartype import beartype +from jax.tree_util import register_pytree_node_class +from typing_extensions import assert_type + +from bearshape import N +from bearshape.jax import Tree as JaxTree +from bearshape.numpy import F32 +from bearshape.optree import Tree + +if TYPE_CHECKING: + from numpy.typing import NDArray + + +@beartype +def optree_int(value: Tree[int]) -> Tree[int]: + return value + + +@beartype +def jax_int(value: JaxTree[int]) -> JaxTree[int]: + return value + + +@beartype +def optree_array(value: Tree[F32[N]]) -> Tree[F32[N]]: + return value + + +@beartype +def jax_array(value: JaxTree[F32[N]]) -> JaxTree[F32[N]]: + return value + + +@beartype +def string_leaf(value: Tree[str]) -> Tree[str]: + return value + + +class Pair(NamedTuple): + first: int + second: int + + +items: list[int] = [1, 2] +nested: list[list[int]] = [[1, 2]] +mapping: dict[str, list[int]] = {"x": items} +values: tuple[int, list[int]] = (1, items) +ordered: OrderedDict[str, int] = OrderedDict(x=1) +empty_list: list[int] = [] +empty_dict: dict[str, int] = {} +array: NDArray[np.float32] = np.ones(3, dtype=np.float32) +arrays: dict[str, list[NDArray[np.float32]]] = {"x": [array, array]} + +optree_int(1) +optree_int(items) +optree_int(nested) +optree_int(mapping) +optree_int(values) +optree_int(ordered) +optree_int(Pair(1, 2)) +optree_int(None) +optree_int(empty_list) +optree_int(empty_dict) +optree_int(()) +assert_type(optree_array(arrays), Tree[F32[N]]) +optree_array(array) +string_leaf("a leaf") +string_leaf(["first", "second"]) +jax_int(1) +jax_int(items) +jax_int(nested) +jax_int(mapping) +jax_int(values) +jax_int(ordered) +jax_int(Pair(1, 2)) +jax_int(None) +jax_int(empty_list) +jax_int(empty_dict) +jax_int(()) +assert_type(jax_array(arrays), JaxTree[F32[N]]) +jax_array(array) + +if TYPE_CHECKING: + assert_type(optree_int(items), Tree[int]) + assert_type(jax_int(items), JaxTree[int]) + +# A registry cannot be inferred statically. Use the existing conditional-alias +# pattern to keep a custom node's concrete type and its runtime leaf check. + + +@register_pytree_node_class +class Batch: + def __init__(self, data: NDArray[np.float32]) -> None: + self.data = data + + def tree_flatten(self) -> tuple[tuple[NDArray[np.float32]], None]: + return (self.data,), None + + @classmethod + def tree_unflatten( + cls, _auxiliary: None, children: tuple[NDArray[np.float32]] + ) -> Batch: + return cls(children[0]) + + +if TYPE_CHECKING: + BatchTree: TypeAlias = Batch +else: + BatchTree = JaxTree[F32[N]] + + +@beartype +def custom_node(value: BatchTree) -> BatchTree: + return value + + +assert_type(custom_node(Batch(array)), Batch) + +assert optree_int(items) is items +assert items == [1, 2] diff --git a/tests/typing_negative/invalid_tree_calls.py b/tests/typing_negative/invalid_tree_calls.py new file mode 100644 index 0000000..4aedc2c --- /dev/null +++ b/tests/typing_negative/invalid_tree_calls.py @@ -0,0 +1,52 @@ +"""Wrong leaves must not disappear inside recursive Tree types.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from bearshape import N +from bearshape.jax import Tree as JaxTree +from bearshape.numpy import F32 +from bearshape.optree import Tree + +if TYPE_CHECKING: + from numpy.typing import NDArray + + +def optree_int(value: Tree[int]) -> None: + pass + + +def jax_int(value: JaxTree[int]) -> None: + pass + + +def optree_array(value: Tree[F32[N]]) -> None: + pass + + +def jax_array(value: JaxTree[F32[N]]) -> None: + pass + + +strings: list[str] = ["wrong"] +nested: dict[str, list[str]] = {"x": strings} +wrong_tuple: tuple[int, str] = (1, "wrong") +wrong_array: NDArray[np.int32] = np.ones(3, dtype=np.int32) +arrays: list[NDArray[np.int32]] = [wrong_array] +optree_int("wrong") # expect: argument +optree_int(strings) # expect: argument +optree_int(nested) # expect: argument +optree_int(wrong_tuple) # expect: argument +optree_int({1, 2}) # expect: argument +optree_array(wrong_array) # expect: argument +optree_array(arrays) # expect: argument +jax_int("wrong") # expect: argument +jax_int(strings) # expect: argument +jax_int(nested) # expect: argument +jax_int(wrong_tuple) # expect: argument +jax_int({1, 2}) # expect: argument +jax_array(wrong_array) # expect: argument +jax_array(arrays) # expect: argument