Better traceback

Author: crusaderky

src/array_api_extra/testing.py

@@ -203,19 +203,35 @@ class CountingDaskScheduler(SchedulerGetCallable):
     """
     Dask scheduler that counts how many times `dask.compute` is called.
 
+    If the number of times exceeds 'max_count', it raises an error.
     This is a wrapper around Dask's own 'synchronous' scheduler.
+
+    Parameters
+    ----------
+    max_count : int
+        Maximum number of allowed calls to `dask.compute`.
+    msg : str
+        Assertion to raise when the count exceeds `max_count`.
     """
 
     count: int
+    max_count: int
+    msg: str
 
-    def __init__(self):  # numpydoc ignore=GL08
+    def __init__(self, max_count: int, msg: str):  # numpydoc ignore=GL08
         self.count = 0
+        self.max_count = max_count
+        self.msg = msg
 
     @override
     def __call__(self, dsk: Graph, keys: Sequence[Key] | Key, **kwargs: Any) -> Any:  # type: ignore[no-any-decorated,no-any-explicit] # numpydoc ignore=GL08
         import dask
 
         self.count += 1
+        # This should yield a nice traceback to the
+        # offending line in the user's code
+        assert self.count <= self.max_count, self.msg
+
         return dask.get(dsk, keys, **kwargs)  # type: ignore[attr-defined,no-untyped-call] # pyright: ignore[reportPrivateImportUsage]
 
 
@@ -228,21 +244,18 @@ def _allow_dask_compute(
     import dask.config
 
     func_name = getattr(func, "__name__", str(func))
+    n_str = f"only up to {n}" if n else "no"
+    msg = (
+        f"Called `dask.compute()` or `dask.persist()` {n + 1} times, "
+        f"but {n_str} calls are allowed. Set "
+        f"`lazy_xp_function({func_name}, allow_dask_compute={n + 1}) "
+        "to allow for more (but note that this will harm performance). "
+    )
 
     @wraps(func)
     def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
-        scheduler = CountingDaskScheduler()
+        scheduler = CountingDaskScheduler(n, msg)
         with dask.config.set({"scheduler": scheduler}):
-            out = func(*args, **kwargs)
-        if scheduler.count > n:
-            n_str = f"only up to {n}" if n else "no"
-            msg = (
-                f"Called `dask.compute()` or `dask.persist()` {scheduler.count} times, "
-                f"but {n_str} calls are allowed. Set "
-                f"`lazy_xp_function({func_name}, allow_dask_compute={scheduler.count}) "
-                "to allow for more (but note that this will harm performance). "
-            )
-            raise AssertionError(msg)
-        return out
+            return func(*args, **kwargs)
 
     return wrapper

tests/test_testing.py

@@ -160,7 +160,7 @@ def test_lazy_xp_function(xp: ModuleType):
             non_materializable3(x)
         with pytest.raises(
             AssertionError,
-            match=r"dask\.compute.* 2 times, but no calls are allowed",
+            match=r"dask\.compute.* 1 times, but no calls are allowed",
         ):
             non_materializable4(x)