Now specifying module and qualname in error messages, for easier debu…

…gging.
patrick-kidger · Jun 21, 2024 · 4b754b8 · 4b754b8
1 parent 075f66e
commit 4b754b8
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 11 deletions.
diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
@@ -438,14 +438,16 @@ def wrapped_fn(*args, **kwargs):
                             typechecker,
                         )
                         try:
-                            name = fn.__name__
+                            module_name = fn.__module__
+                            qualname = fn.__qualname__
                         except AttributeError:
-                            name = fn.__class__.__name__
+                            module_name = fn.__class__.__module__
+                            qualname = fn.__class__.__qualname__
                         param_values = _pformat(bound.arguments, short_self=True)
                         param_hints = _remove_typing(param_signature)
                         msg = (
                             "Type-check error whilst checking the parameters of "
-                            f"{name}.{argmsg}\n"
+                            f"{module_name}.{qualname}.{argmsg}\n"
                             "----------------------\n"
                             f"Called with parameters: {param_values}\n"
                             f"Parameter annotations: {param_hints}.\n"
@@ -481,9 +483,11 @@ def wrapped_fn(*args, **kwargs):
                             raise
                         except Exception as e:
                             try:
-                                name = fn.__name__
+                                module_name = fn.__module__
+                                qualname = fn.__qualname__
                             except AttributeError:
-                                name = fn.__class__.__name__
+                                module_name = fn.__class__.__module__
+                                qualname = fn.__class__.__qualname__
                             param_values = _pformat(bound.arguments, short_self=True)
                             return_value = _pformat(out, short_self=False)
                             param_hints = _remove_typing(param_signature)
@@ -496,7 +500,7 @@ def wrapped_fn(*args, **kwargs):
                                 return_hint = return_hint[8:-2]
                             msg = (
                                 "Type-check error whilst checking the return value "
-                                f"of {name}.\n"
+                                f"of {module_name}.{qualname}.\n"
                                 f"Actual value: {return_value}\n"
                                 f"Expected type: {return_hint}.\n"
                                 "----------------------\n"
@@ -565,6 +569,7 @@ def _check_dataclass_annotations(self, typechecker):
     f = _make_fn_with_signature(
         self.__class__.__name__, signature, module, output=False
     )
+    f.__qualname__ = self.__class__.__qualname__
     f = jaxtyped(f, typechecker=typechecker)
     f(self, **values)
 

diff --git a/test/test_generators.py b/test/test_generators.py
@@ -1,4 +1,4 @@
-from typing import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator
 
 import jax.numpy as jnp
 import pytest

diff --git a/test/test_messages.py b/test/test_messages.py
@@ -13,7 +13,7 @@ def f(x: str, y: str, z: int):
         pass
 
     matches = [
-        "Type-check error whilst checking the parameters of f",
+        "Type-check error whilst checking the parameters of .*<locals>.f",
         "The problem arose whilst typechecking parameter 'z'.",
         "Called with parameters: {'x': 'hi', 'y': 'bye', 'z': 'not-an-int'}",
         r"Parameter annotations: \(x: str, y: str, z: int\).",
@@ -29,7 +29,7 @@ def g(x: Float[Array, "a b"], y: Float[Array, "b c"]):
     x = jnp.zeros((2, 3))
     y = jnp.zeros((4, 3))
     matches = [
-        "Type-check error whilst checking the parameters of g",
+        "Type-check error whilst checking the parameters of .*.<locals>.g",
         "The problem arose whilst typechecking parameter 'y'.",
         r"Called with parameters: {'x': f32\[2,3\], 'y': f32\[4,3\]}",
         (
@@ -53,7 +53,7 @@ def f(x: PyTree[Any, " T"], y: PyTree[Any, " S"]) -> PyTree[Any, "T S"]:
     x = (1, 2)
     y = {"a": 1}
     matches = [
-        "Type-check error whilst checking the return value of f",
+        "Type-check error whilst checking the return value of .*.<locals>.f",
         r"Called with parameters: {'x': \(1, 2\), 'y': {'a': 1}}",
         "Actual value: 'foo'",
         r"Expected type: PyTree\[Any, \"T S\"\].",
@@ -81,7 +81,7 @@ class M(eqx.Module):
     z = "not-an-int"
 
     matches = [
-        "Type-check error whilst checking the parameters of M",
+        "Type-check error whilst checking the parameters of .*.<locals>.M",
         "The problem arose whilst typechecking parameter 'z'.",
         (
             r"Called with parameters: {'self': M\(\.\.\.\), 'x': f32\[2,3\], "