Pre-commit hooks for pyright and mypy (only API tests) (#283)

* Added pre-commit hooks for pyright and mypy (only API tests)

* Removed unnecessary noqa

* Replaced tensorflow for numpy in type tests

* Try tweaking pre-commit-config

---------

Co-authored-by: Martí Zamora Casals <[email protected]>
Co-authored-by: Patrick Kidger <[email protected]>

.pre-commit-config.yaml

@@ -26,3 +26,18 @@ repos:
       - id: ruff  # linter
         types_or: [ python, pyi, jupyter ]
         args: [ --fix ]
+  - repo: https://github.com/RobertCraigie/pyright-python
+    rev: v1.1.391
+    hooks:
+    - id: pyright
+      files: ^test/types/
+      additional_dependencies:
+          [beartype, numpy<2]
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.14.1  
+    hooks:
+      - id: mypy
+        files: ^test/types/
+        additional_dependencies:
+          [beartype, numpy<2]
+        args: ["--ignore-missing-imports", "--follow-imports=skip"]

test/types/decorator.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+
+import numpy as np
+from beartype import beartype
+
+from jaxtyping import Float, Int, jaxtyped
+
+
+@jaxtyped(typechecker=beartype)
+@dataclass
+class User:
+    name: str
+    age: int
+    items: Float[np.ndarray, " N"]
+    timestamps: Int[np.ndarray, " N"]
+
+
+@jaxtyped(typechecker=beartype)
+def transform_user(user: User, increment_age: int = 1) -> User:
+    user.age += increment_age
+    return user
+
+
+user = User(
+    name="John",
+    age=20,
+    items=np.random.normal(size=10),
+    timestamps=np.random.randint(0, 100, size=10),
+)
+
+new_user = transform_user(user, increment_age=2)