Fix memo leaking to nested beartype helpers, add TreeChecker replay guard, document Tree structure syntax as runtime-only

- Tag explicit memo entries with owner_code so @shapix.check scope does
  not bleed into nested plain @beartype calls (Finding 1)
- Add _fail_obj replay guard and memo snapshot/restore to _TreeChecker
  to prevent contradictory beartype error messages (Finding 3)
- Document Tree structure args (T, S, ...) as runtime-only in module
  docstring, README, and typing fixture (Finding 2)

Author: acecchini

README.md

@@ -799,6 +799,7 @@ However, some patterns are fundamentally runtime-only and produce type checker e
 | Arithmetic | `F32[N + 2]` | `# type: ignore` |
 | Custom dimensions | `F32[Vocab, Embed]` | `# type: ignore` or `TYPE_CHECKING` pattern |
 | `Value(...)` | `F32[Value("size")]` | `# type: ignore` |
+| Tree structure args | `Tree[F32[N], T]`, `Tree[F32[N], T, ...]` | `# type: ignore` — leaf-only `Tree[F32[N, C]]` works |
 
 ### Recommended type checker config
 

src/shapix/_decorator.py

@@ -59,24 +59,27 @@ def decorator(fn: Callable[P, R]) -> Callable[P, R]:
       inner = beartype(fn, conf=conf)  # type: ignore[call-overload]
 
     if inspect.iscoroutinefunction(fn):
+      inner_code = getattr(inner, "__code__", None)
 
       @functools.wraps(fn)
       async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
         bound = signature.bind_partial(*args, **kwargs)
         bound.apply_defaults()
-        push_memo(scope=dict(bound.arguments))
+        push_memo(scope=dict(bound.arguments), owner_code=inner_code)
         try:
           return await inner(*args, **kwargs)  # type: ignore[misc,no-any-return]
         finally:
           pop_memo()
 
       return async_wrapper  # type: ignore[return-value]
 
+    inner_code = getattr(inner, "__code__", None)
+
     @functools.wraps(fn)
     def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
       bound = signature.bind_partial(*args, **kwargs)
       bound.apply_defaults()
-      push_memo(scope=dict(bound.arguments))
+      push_memo(scope=dict(bound.arguments), owner_code=inner_code)
       try:
         return inner(*args, **kwargs)
       finally:

src/shapix/_memo.py

@@ -79,23 +79,39 @@ def restore(
 _explicit_scope_stack: contextvars.ContextVar[tuple[dict[str, object] | None, ...]] = (
   contextvars.ContextVar("_explicit_scope_stack", default=())
 )
+_explicit_owner_stack: contextvars.ContextVar[tuple[types.CodeType | None, ...]] = (
+  contextvars.ContextVar("_explicit_owner_stack", default=())
+)
 
 
 def push_memo(
-  memo: ShapeMemo | None = None, *, scope: dict[str, object] | None = None
+  memo: ShapeMemo | None = None,
+  *,
+  scope: dict[str, object] | None = None,
+  owner_code: types.CodeType | None = None,
 ) -> ShapeMemo:
-  """Push a memo onto the explicit stack. Pair with :func:`pop_memo`."""
+  """Push a memo onto the explicit stack. Pair with :func:`pop_memo`.
+
+  Parameters
+  ----------
+  owner_code:
+      When set, this entry is only visible to frames whose ``f_code``
+      matches.  Untagged entries (``None``) are unconditional and used
+      by :class:`check_context` and :class:`_TreeChecker`.
+  """
   if memo is None:
     memo = ShapeMemo()
   _explicit_stack.set(_explicit_stack.get() + (memo,))
   _explicit_scope_stack.set(_explicit_scope_stack.get() + (scope,))
+  _explicit_owner_stack.set(_explicit_owner_stack.get() + (owner_code,))
   return memo
 
 
 def pop_memo() -> None:
   """Pop the most recent explicit memo."""
   _explicit_stack.set(_explicit_stack.get()[:-1])
   _explicit_scope_stack.set(_explicit_scope_stack.get()[:-1])
+  _explicit_owner_stack.set(_explicit_owner_stack.get()[:-1])
 
 
 # ---------------------------------------------------------------------------
@@ -122,10 +138,22 @@ def get_memo(_depth: int = 2) -> ShapeMemo:
       Default ``2`` accounts for: our validator → beartype's ``_is_valid_bool``
       → beartype wrapper.
   """
-  # 1. Explicit stack takes priority
+  # 1. Explicit stack takes priority (if owner matches or untagged)
   explicit = _explicit_stack.get()
   if explicit:
-    return explicit[-1]
+    owners = _explicit_owner_stack.get()
+    owner_code = owners[-1] if owners else None
+    if owner_code is None:
+      # Untagged entry (check_context / TreeChecker) — always visible
+      return explicit[-1]
+    # Tagged entry — only visible if the caller's frame matches
+    try:
+      frame = sys._getframe(_depth)
+      if frame.f_code is owner_code:
+        return explicit[-1]
+    except ValueError:
+      pass
+    # Fall through to frame-based detection
 
   # 2. Frame-based detection
   try:
@@ -187,7 +215,19 @@ def get_scope(_depth: int = 2) -> dict[str, object]:
   """Return the current runtime scope for ``Value(...)`` expressions."""
   explicit = _explicit_scope_stack.get()
   if explicit and explicit[-1] is not None:
-    return explicit[-1]
+    owners = _explicit_owner_stack.get()
+    owner_code = owners[-1] if owners else None
+    if owner_code is None:
+      # Untagged entry — always visible
+      return explicit[-1]
+    # Tagged entry — only visible if the caller's frame matches
+    try:
+      frame = sys._getframe(_depth)
+      if frame.f_code is owner_code:
+        return explicit[-1]
+    except ValueError:
+      pass
+    # Fall through to frame-based detection
 
   try:
     frame = sys._getframe(_depth)

src/shapix/_tree.py

@@ -9,6 +9,14 @@
 Requires ``optree`` or ``jax`` for tree traversal. Install with
 ``pip install optree`` or ``pip install jax``.
 
+.. note::
+
+   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.
+
 Import ``Tree`` from an explicit backend module::
 
     from shapix.optree import Tree  # backed by optree
@@ -86,7 +94,7 @@ def __repr__(self) -> str:
 class _TreeChecker:
   """Beartype validator for tree leaf types and structure consistency."""
 
-  __slots__ = ("_leaf_type", "_structure_spec", "_get_ops", "_repr")
+  __slots__ = ("_leaf_type", "_structure_spec", "_get_ops", "_repr", "_fail_obj")
 
   def __init__(
     self,
@@ -98,10 +106,20 @@ def __init__(
     self._leaf_type = leaf_type
     self._structure_spec = structure_spec
     self._get_ops = get_ops
+    self._fail_obj: object | None = None
     spec_str = f", {structure_spec}" if structure_spec else ""
     self._repr = f"Tree[{leaf_type!r}{spec_str}]"
 
   def __call__(self, obj: object) -> bool:
+    # Replay guard: when beartype's error-generation code re-invokes us
+    # from a different call-stack frame, the fresh memo would lack prior
+    # bindings, causing a previously failing check to pass.  Keep replaying
+    # the failure until a successful validation with a (possibly different)
+    # object clears it.  We use ``is`` identity (not ``id()``) to avoid
+    # false matches from recycled addresses.
+    if self._fail_obj is not None and self._fail_obj is obj:
+      return False
+
     tree_ops = self._get_ops()
     from beartype.door import is_bearable
 
@@ -111,12 +129,21 @@ def __call__(self, obj: object) -> bool:
     # and can resolve ``Value(...)`` expressions against the same parameters.
     memo = get_memo(_depth=3)
     scope = get_scope(_depth=3)
+    snap = memo.snapshot()
     push_memo(memo, scope=scope)
     try:
-      return self._validate(obj, tree_ops, is_bearable, memo)
+      result = self._validate(obj, tree_ops, is_bearable, memo)
     finally:
       pop_memo()
 
+    if not result:
+      memo.restore(snap)
+      self._fail_obj = obj
+    else:
+      self._fail_obj = None
+
+    return result
+
   def _validate(
     self, obj: object, tree_ops: tp.Any, is_bearable: tp.Any, memo: ShapeMemo
   ) -> bool:

tests/test_decorator.py

@@ -208,6 +208,89 @@ def f(size: int) -> F32[Value("size")]:  # type: ignore[valid-type]
     assert f(7).shape == (7,)
 
 
+class TestMemoIsolation:
+  """Explicit memo from @shapix.check must not leak to nested @beartype helpers."""
+
+  def test_check_outer_plain_beartype_inner_independent_dims(self) -> None:
+    """Outer @shapix.check binds N=4, inner plain @beartype binds N=7."""
+
+    @beartype
+    def inner(x: F32[N]) -> F32[N]:
+      return x
+
+    @shapix.check
+    @beartype
+    def outer(x: F32[N]) -> F32[N]:
+      # Call inner with a different N — must succeed independently
+      inner(np.ones(7, dtype=np.float32))
+      return x
+
+    outer(np.ones(4, dtype=np.float32))
+
+  def test_check_outer_plain_beartype_inner_value(self) -> None:
+    """Inner plain @beartype with Value resolves from its own frame."""
+
+    @beartype
+    def inner(size: int) -> F32[Value("size")]:  # type: ignore[valid-type]
+      return np.ones(size, dtype=np.float32)
+
+    @shapix.check
+    @beartype
+    def outer(dummy: int) -> F32[Value("dummy")]:  # type: ignore[valid-type]
+      inner(7)
+      return np.ones(dummy, dtype=np.float32)
+
+    outer(4)
+
+  def test_check_context_shares_memo_with_beartype_call(self) -> None:
+    """Inside check_context(), @beartype calls share the same memo (by design)."""
+
+    @beartype
+    def helper(x: F32[N]) -> F32[N]:
+      return x
+
+    with shapix.check_context():
+      from beartype.door import is_bearable
+
+      assert is_bearable(np.ones(4, dtype=np.float32), F32[N])
+      # check_context is untagged — helper shares the memo, so N=4 is bound
+      helper(np.ones(4, dtype=np.float32))  # same N — OK
+      with pytest.raises(BeartypeCallHintParamViolation):
+        helper(np.ones(7, dtype=np.float32))  # different N — fails
+
+  def test_async_check_outer_plain_beartype_inner(self) -> None:
+    """Async variant: outer @shapix.check does not leak to inner @beartype."""
+    import asyncio
+
+    @beartype
+    def inner(x: F32[N]) -> F32[N]:
+      return x
+
+    @shapix.check
+    @beartype
+    async def outer(x: F32[N]) -> F32[N]:
+      inner(np.ones(7, dtype=np.float32))
+      return x
+
+    asyncio.run(outer(np.ones(4, dtype=np.float32)))
+
+  def test_async_child_task_independent_memo(self) -> None:
+    """Child task spawned inside @shapix.check gets isolated memo after parent returns."""
+    import asyncio
+
+    @shapix.check
+    @beartype
+    async def parent(x: F32[N]) -> F32[N]:
+      return x
+
+    async def run() -> None:
+      await parent(np.ones(4, dtype=np.float32))
+      # After parent returns, a new call with different N works
+      await parent(np.ones(10, dtype=np.float32))
+
+    asyncio.run(run())
+
+
 class TestAsyncCheckContext:
   def test_async_check_context(self) -> None:
     """async with check_context() works correctly."""

tests/test_tree.py

@@ -1300,6 +1300,44 @@ def test_parse_spec_args_empty_tuple_error(self) -> None:
       Tree[()]
 
 
+# =====================================================================
+# Replay guard (_fail_obj) — non-contradictory error messages
+# =====================================================================
+
+
+class TestReplayGuard:
+  """_TreeChecker must produce non-contradictory beartype error messages."""
+
+  def test_plain_beartype_tree_param_mismatch_error_message(self) -> None:
+    """Plain @beartype tree param mismatch must not say True == Is[...]."""
+
+    @shapix.check
+    @beartype
+    def f(x: Tree[F32[N], T], y: Tree[F32[N], T]) -> Tree[F32[N]]:
+      return x
+
+    x = {"a": np.ones(3, dtype=np.float32)}
+    y = [np.ones(3, dtype=np.float32)]  # different structure
+    with pytest.raises(BeartypeCallHintParamViolation) as exc_info:
+      f(x, y)
+    # The error message should NOT contain "True ==" which would indicate
+    # the replay guard failed (validator passed on re-invocation)
+    assert "True ==" not in str(exc_info.value)
+
+  def test_plain_beartype_tree_return_mismatch_error_message(self) -> None:
+    """Return type tree mismatch must produce non-contradictory error."""
+
+    @shapix.check
+    @beartype
+    def f(x: Tree[F32[N], T]) -> Tree[F32[N], T]:
+      # Return different structure than input
+      return [np.ones(3, dtype=np.float32)]
+
+    with pytest.raises(BeartypeCallHintReturnViolation) as exc_info:
+      f({"a": np.ones(3, dtype=np.float32)})
+    assert "True ==" not in str(exc_info.value)
+
+
 # =====================================================================
 # Backend-specific Tree modules
 # =====================================================================

tests/typing/check_tree.py

@@ -67,3 +67,16 @@ def identity(x: Tree[object]) -> Tree[object]:
 
 def to_none(x: Tree[object]) -> None:
   pass
+
+
+# ---------------------------------------------------------------------------
+# Note: Structure args (T, S, ...) are runtime-only
+# ---------------------------------------------------------------------------
+# Multi-arg syntax like Tree[F32[N], T] or Tree[F32[N], T, ...] works at
+# runtime but is not understood by type checkers (Tree has one type param).
+# Use `# type: ignore` for structure-bearing annotations.
+# Leaf-only annotations like Tree[int] or Tree[object] work with all checkers.
+
+
+def leaf_int(x: Tree[int]) -> Tree[int]:
+  return x