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

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.",