From 5fbd6718ab5681d99cdfcfbe19c5e1e1042d9e49 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Thu, 2 Nov 2023 20:01:05 -0700 Subject: [PATCH] Added support for 'self' in dataclass attribute annotations; switched from args and kwargs to just arguments. --- jaxtyping/_decorator.py | 46 +++++++++++++++++++++++++---------------- test/test_array.py | 15 ++++++++++++++ test/test_messages.py | 16 ++++++-------- 3 files changed, 49 insertions(+), 28 deletions(-) diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index dd57a72..d16b1ab 100644 --- a/jaxtyping/_decorator.py +++ b/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( ""): - 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 diff --git a/test/test_array.py b/test/test_array.py index 3be95f1..1b81cae 100644 --- a/test/test_array.py +++ b/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, ""] diff --git a/test/test_messages.py b/test/test_messages.py index 2d8f0ad..2568baf 100644 --- a/test/test_messages.py +++ b/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.",