Call typechecker from isolated frame (#260)

* Call typechecker from isolated frame

* Add comment

* Enable test_no_garbage for TypeGuard

jaxtyping/_decorator.py

@@ -57,6 +57,15 @@ def __repr__(self):
 _tb_flag = True
 
 
+def _apply_typechecker(typechecker, fn):
+    """Calls `typechecker(fn)` in an isolated frame, returning the result.
+
+    This avoids reference cycles that can otherwise occur if `typechecker` grabs
+    the calling frame's locals.
+    """
+    return typechecker(fn)
+
+
 @overload
 def jaxtyped(
     *,
@@ -422,8 +431,8 @@ def wrapped_fn(*args, **kwargs):  # pyright: ignore
             param_fn = _make_fn_with_signature(
                 name, qualname, module, param_signature, output=False
             )
-            full_fn = typechecker(full_fn)
-            param_fn = typechecker(param_fn)
+            full_fn = _apply_typechecker(typechecker, full_fn)
+            param_fn = _apply_typechecker(typechecker, param_fn)
 
             def wrapped_fn_impl(args, kwargs, bound, memos):
                 # First type-check just the parameters before the function is
@@ -790,7 +799,9 @@ def _get_problem_arg(
         fn = _make_fn_with_signature(
             "check_single_arg", "check_single_arg", module, new_signature, output=False
         )
-        fn = typechecker(fn)  # but no `jaxtyped`; keep the same environment.
+        fn = _apply_typechecker(
+            typechecker, fn
+        )  # but no `jaxtyped`; keep the same environment.
         try:
             fn(*args, **kwargs)
         except Exception as e:

test/test_decorator.py

@@ -1,11 +1,11 @@
 import abc
 import dataclasses
+import sys
 from typing import no_type_check
 
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
-import typeguard
 
 from jaxtyping import Array, Float, jaxtyped, print_bindings
 
@@ -213,10 +213,6 @@ def g(x: Float[Array, "foo bar"]):
 
 
 def test_no_garbage(typecheck):
-    if typecheck is typeguard.typechecked:
-        # Currently fails due to reference cycles in typeguard.
-        pytest.skip()
-
     with assert_no_garbage():
 
         @jaxtyped(typechecker=typecheck)
@@ -236,3 +232,28 @@ class _Obj:
             x: int
 
         _Obj(x=5)
+
+
+def test_no_garbage_frame_capture_typecheck():
+    with assert_no_garbage():
+        # Some typechecker implementations (e.g., typeguard 2.13.3) capture the calling
+        # frame's f_locals. This test checks that the calling frames in jaxtyping are
+        # sufficiently isolated to avoid introducing reference cycles when a
+        # typechecker does this.
+        def frame_locals_capture(fn):
+            locals = sys._getframe(1).f_locals
+
+            def wrapper(*args, **kwargs):
+                # Required to ensure wrapper holds a reference to f_locals, which is
+                # the scenario under test.
+                _ = locals
+                return fn(*args, **kwargs)
+
+            return wrapper
+
+        @jaxtyped(typechecker=frame_locals_capture)
+        @dataclasses.dataclass
+        class _Obj:
+            x: int
+
+        _Obj(x=5)