Skip to content

Commit

Permalink
Now specifying module and qualname in error messages, for easier debu…
Browse files Browse the repository at this point in the history
…gging.
  • Loading branch information
patrick-kidger committed Jun 21, 2024
1 parent 075f66e commit 4b754b8
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 11 deletions.
17 changes: 11 additions & 6 deletions jaxtyping/_decorator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
Expand All @@ -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"
Expand Down Expand Up @@ -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)

Expand Down
2 changes: 1 addition & 1 deletion test/test_generators.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import AsyncIterator, Iterator
from collections.abc import AsyncIterator, Iterator

import jax.numpy as jnp
import pytest
Expand Down
8 changes: 4 additions & 4 deletions test/test_messages.py
Original file line number Diff line number Diff line change
Expand Up @@ -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\).",
Expand All @@ -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\]}",
(
Expand All @@ -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\"\].",
Expand Down Expand Up @@ -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\], "
Expand Down

0 comments on commit 4b754b8

Please sign in to comment.