diff --git a/jaxtyping/import_hook.py b/jaxtyping/import_hook.py index 6d66e76..a1ed323 100644 --- a/jaxtyping/import_hook.py +++ b/jaxtyping/import_hook.py @@ -111,7 +111,7 @@ class _JaxtypingTransformer(ast.NodeVisitor): args = [ast.Constant(None)] else: args = [_dot_lookup(*self._typechecker)] - node.decorator_list.append(ast.Call(func, args, keywords=[])) + node.decorator_list.insert(0, ast.Call(func, args, keywords=[])) self._parents.append(node) self.generic_visit(node) self._parents.pop() diff --git a/test/import_hook_tester_beartype.py b/test/import_hook_tester_beartype.py index db1cd29..b77b756 100644 --- a/test/import_hook_tester_beartype.py +++ b/test/import_hook_tester_beartype.py @@ -17,6 +17,8 @@ # 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 + import equinox as eqx import jax.numpy as jnp import pytest @@ -45,3 +47,16 @@ with pytest.raises(ParamError): M(1.0, jnp.array([1.0])) with pytest.raises(ParamError): M(1, jnp.array(1.0)) + + +@dataclasses.dataclass +class D: + foo: int + bar: Float32[jnp.ndarray, " a"] + + +D(1, jnp.array([1.0])) +with pytest.raises(ParamError): + D(1.0, jnp.array([1.0])) +with pytest.raises(ParamError): + D(1, jnp.array(1.0)) diff --git a/test/import_hook_tester_broken_checker.py b/test/import_hook_tester_broken_checker.py index db1cd29..b77b756 100644 --- a/test/import_hook_tester_broken_checker.py +++ b/test/import_hook_tester_broken_checker.py @@ -17,6 +17,8 @@ # 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 + import equinox as eqx import jax.numpy as jnp import pytest @@ -45,3 +47,16 @@ with pytest.raises(ParamError): M(1.0, jnp.array([1.0])) with pytest.raises(ParamError): M(1, jnp.array(1.0)) + + +@dataclasses.dataclass +class D: + foo: int + bar: Float32[jnp.ndarray, " a"] + + +D(1, jnp.array([1.0])) +with pytest.raises(ParamError): + D(1.0, jnp.array([1.0])) +with pytest.raises(ParamError): + D(1, jnp.array(1.0)) diff --git a/test/import_hook_tester_typeguard.py b/test/import_hook_tester_typeguard.py index db1cd29..b77b756 100644 --- a/test/import_hook_tester_typeguard.py +++ b/test/import_hook_tester_typeguard.py @@ -17,6 +17,8 @@ # 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 + import equinox as eqx import jax.numpy as jnp import pytest @@ -45,3 +47,16 @@ with pytest.raises(ParamError): M(1.0, jnp.array([1.0])) with pytest.raises(ParamError): M(1, jnp.array(1.0)) + + +@dataclasses.dataclass +class D: + foo: int + bar: Float32[jnp.ndarray, " a"] + + +D(1, jnp.array([1.0])) +with pytest.raises(ParamError): + D(1.0, jnp.array([1.0])) +with pytest.raises(ParamError): + D(1, jnp.array(1.0))