Fix vanilla dataclasses (#56)

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()

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))

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))

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))