Added support for 'self' in dataclass attribute annotations; switched from args and kwargs to just arguments.

jaxtyping/_decorator.py

@@ -342,13 +342,13 @@ def jaxtyped(fn=None, *, typechecker=None):
                                 name = fn.__name__
                             except AttributeError:
                                 name = fn.__class__.__name__
-                            paramstr = _remove_typing(param_signature)
+                            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"Called with args: {_pformat(args)}\n"
-                                f"Called with kwargs: {_pformat(kwargs)}\n"
-                                f"Parameter annotations: {paramstr}.\n"
+                                f"Called with arguments: {param_values}\n"
+                                f"Parameter annotations: {param_hints}.\n"
                                 + _exc_shape_info(memos)
                             )
                             raise TypeCheckError(msg) from e
@@ -382,23 +382,25 @@ def jaxtyped(fn=None, *, typechecker=None):
                                     name = fn.__name__
                                 except AttributeError:
                                     name = fn.__class__.__name__
-                                paramstr = _remove_typing(param_signature)
-                                returnstr = _remove_typing(
+                                param_values = _pformat(
+                                    bound.arguments, short_self=True
+                                )
+                                return_value = _pformat(out, short_self=False)
+                                param_hints = _remove_typing(param_signature)
+                                return_hint = _remove_typing(
                                     full_signature.return_annotation
                                 )
-                                if returnstr.startswith(
+                                if return_hint.startswith(
                                     "<class '"
-                                ) and returnstr.endswith("'>"):
-                                    returnstr = returnstr[8:-2]
-                                kwargs.pop(output_name)
+                                ) and return_hint.endswith("'>"):
+                                    return_hint = return_hint[8:-2]
                                 msg = (
                                     "Type-check error whilst checking the return value "
                                     f"of {name}.\n"
-                                    f"Called with args: {_pformat(args)}\n"
-                                    f"Called with kwargs: {_pformat(kwargs)}\n"
-                                    f"Return value: {_pformat(out)}\n"
-                                    f"Parameter annotations: {paramstr}.\n"
-                                    f"Return annotation: {returnstr}.\n"
+                                    f"Called with arguments: {param_values}\n"
+                                    f"Return value: {return_value}\n"
+                                    f"Parameter annotations: {param_hints}.\n"
+                                    f"Return annotation: {return_hint}.\n"
                                     + _exc_shape_info(memos)
                                 )
                                 raise TypeCheckError(msg) from e
@@ -425,7 +427,7 @@ def _check_dataclass_annotations(self, typechecker):
     `self` should be a dataclass instancae. `typechecker` should be e.g.
     `beartype.beartype` or `typeguard.typechecked`.
     """
-    parameters = []
+    parameters = [inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)]
     values = {}
     for field in dataclasses.fields(self):
         annotation = field.type
@@ -461,7 +463,7 @@ def _check_dataclass_annotations(self, typechecker):
         self.__class__.__name__, signature, module, output=False
     )
     f = jaxtyped(f, typechecker=typechecker)
-    f(**values)
+    f(self, **values)
 
 
 def _make_fn_with_signature(
@@ -655,7 +657,7 @@ def _remove_typing(x):
     return x
 
 
-def _pformat(x):
+def _pformat(x, short_self: bool):
     # No performance concerns from delayed imports -- this is only used when we're about
     # to raise an error anyway.
     try:
@@ -665,6 +667,14 @@ def _pformat(x):
         import equinox as eqx
 
         pformat = eqx.tree_pformat
+        if short_self:
+            try:
+                self = x["self"]
+            except KeyError:
+                pass
+            else:
+                is_self = lambda y: y is self
+                pformat = ft.partial(pformat, truncate_leaf=is_self)
     except Exception:
         import pprint
 

test/test_array.py

@@ -17,6 +17,7 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import dataclasses as dc
 import sys
 from typing import get_args, get_origin, Union
 
@@ -33,6 +34,7 @@ from jaxtyping import (
     Bool,
     Float,
     Float32,
+    jaxtyped,
     Key,
     PRNGKeyArray,
     Scalar,
@@ -485,6 +487,19 @@ def test_deferred_symbolic_bad(jaxtyp, typecheck):
         A().bar(jnp.array(0.0))
 
 
+def test_deferred_symbolic_dataclass(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    @dc.dataclass
+    class A:
+        value: int
+        array: Float[Array, " {self.value}"]
+
+    A(3, jnp.zeros(3))
+
+    with pytest.raises(ParamError):
+        A(3, jnp.zeros(4))
+
+
 def test_arraylike(typecheck, getkey):
     floatlike1 = Float32[ArrayLike, ""]
     floatlike2 = Float[ArrayLike, ""]

test/test_messages.py

@@ -15,8 +15,7 @@ def test_arg_localisation(typecheck):
     matches = [
         "Type-check error whilst checking the parameters of f",
         "The problem arose whilst typechecking argument 'z'.",
-        r"Called with args: \('hi', 'bye', 'not-an-int'\)",
-        "Called with kwargs: {}",
+        "Called with arguments: {'x': 'hi', 'y': 'bye', 'z': 'not-an-int'}",
         r"Parameter annotations: \(x: str, y: str, z: int\).",
     ]
     for match in matches:
@@ -32,8 +31,7 @@ def test_arg_localisation(typecheck):
     matches = [
         "Type-check error whilst checking the parameters of g",
         "The problem arose whilst typechecking argument 'y'.",
-        r"Called with args: \(f32\[2,3\],\)",
-        r"Called with kwargs: {'y': f32\[4,3\]}",
+        r"Called with arguments: {'x': f32\[2,3\], 'y': f32\[4,3\]}",
         (
             r"Parameter annotations: \(x: Float\[Array, 'a b'\], y: "
             r"Float\[Array, 'b c'\]\)."
@@ -56,8 +54,7 @@ def test_return(typecheck):
     y = {"a": 1}
     matches = [
         "Type-check error whilst checking the return value of f",
-        r"Called with args: \(\(1, 2\),\)",
-        r"Called with kwargs: {'y': {'a': 1}}",
+        r"Called with arguments: {'x': \(1, 2\), 'y': {'a': 1}}",
         "Return value: 'foo'",
         r"Return annotation: PyTree\[Any, \"T S\"\].",
         (
@@ -86,13 +83,12 @@ def test_dataclass_attribute(typecheck):
     matches = [
         "Type-check error whilst checking the parameters of M",
         "The problem arose whilst typechecking argument 'z'.",
-        r"Called with args: \(\)",
         (
-            r"Called with kwargs: {'x': f32\[2,3\], 'y': \(1, \(3, 4\)\), "
-            r"'z': 'not-an-int'}"
+            r"Called with arguments: {'self': M\(\.\.\.\), 'x': f32\[2,3\], "
+            r"'y': \(1, \(3, 4\)\), 'z': 'not-an-int'}"
         ),
         (
-            r"Parameter annotations: \(x: Float\[Array, '\*foo'\], "
+            r"Parameter annotations: \(self: Any, x: Float\[Array, '\*foo'\], "
             r"y: PyTree\[Any, \"T\"\], z: int\)."
         ),
         "The current values for each jaxtyping axis annotation are as follows.",