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

…gging.

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)
 

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

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