fix: model ordinary Tree containers for static consumers

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.
 

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:

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.

plans/2026-09-08-static-tree-containers.md

@@ -0,0 +1,89 @@
+# 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.
+- [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 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. 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
+
+
+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
+
+
+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. 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
+
+
+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.

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"]
 
 

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")

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")

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)))

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
 
 # ---------------------------------------------------------------------------