ci: gate development with four python type checkers

.github/workflows/ci.yml

@@ -15,6 +15,20 @@ concurrency:
   cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 jobs:
+  static:
+    name: Lint and type check
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v10.2.0
+        with:
+          python-version: "3.12"
+          enable-cache: true
+      - name: Install
+        run: uv sync --locked
+      - name: Lint and all four type checkers
+        run: make lint type
+
   check:
     runs-on: ubuntu-latest
     strategy:
@@ -29,8 +43,8 @@ jobs:
           enable-cache: true
       - name: Install
         run: uv sync --locked
-      - name: Check (lint + types + tests + coverage)
-        run: make check
+      - name: Tests and coverage
+        run: make coverage
         env:
           CI: "1"
       - name: Examples and notebooks

.pre-commit-config.yaml

@@ -42,6 +42,35 @@ repos:
         language: system
         files: ^(pyproject\.toml|uv\.lock)$
         pass_filenames: false
+      # Whole-project checks: scope belongs to pyproject.toml, commands to Make.
+      - id: mypy
+        name: mypy (strict)
+        entry: make type-mypy
+        language: system
+        files: \.(py|pyi)$|^(pyproject\.toml|uv\.lock|Makefile|\.pre-commit-config\.yaml)$
+        pass_filenames: false
+        require_serial: true
+      - id: pyright
+        name: pyright
+        entry: make type-pyright
+        language: system
+        files: \.(py|pyi)$|^(pyproject\.toml|uv\.lock|Makefile|\.pre-commit-config\.yaml)$
+        pass_filenames: false
+        require_serial: true
+      - id: zuban
+        name: zuban
+        entry: make type-zuban
+        language: system
+        files: \.(py|pyi)$|^(pyproject\.toml|uv\.lock|Makefile|\.pre-commit-config\.yaml)$
+        pass_filenames: false
+        require_serial: true
+      - id: pyrefly
+        name: pyrefly
+        entry: make type-pyrefly
+        language: system
+        files: \.(py|pyi)$|^(pyproject\.toml|uv\.lock|Makefile|\.pre-commit-config\.yaml)$
+        pass_filenames: false
+        require_serial: true
   - repo: https://github.com/compilerla/conventional-pre-commit
     rev: v4.4.0
     hooks:

AGENTS.md

@@ -12,9 +12,17 @@ Run everything through `make` or `uv run`. The system `python` may be
 older than the required 3.12; uv supplies the pinned interpreter and the
 locked environment (`uv.lock`).
 
-- `make check`: the gate (ruff, strict mypy, pytest, 95% branch coverage),
-  about 10 s. The work is done when it passes.
+- `make check`: the gate (cleanporter, ruff, mypy, pyright, zuban, pyrefly,
+  pytest, 95% branch coverage). The work is done when it passes.
 - `make lint type test`: the same gate without coverage, for the inner loop.
+- `make type`: all four blocking type checkers; `make type-mypy`,
+  `make type-pyright`, `make type-zuban`, `make type-pyrefly` run one.
+  Scope and settings live in `pyproject.toml`: mypy (strict) and pyright
+  (standard) check `src`; zuban (strict, test annotation requirements relaxed)
+  and pyrefly also check tests. Corpus input data, generated golden cases,
+  and optional JAX/PyTorch compilation tests are excluded from the wider pair.
+  Test helpers describe dynamic callable metadata and assert IR binding shapes;
+  use them instead of adding blanket type suppressions.
 - `make fmt`: format and autofix.
 - Area-specific: `make fuzz` (lowering, optimizer, runtime), `make backends`
   (codegen, namespace handling), `make docs` (anything in `docs/`),

Makefile

@@ -12,7 +12,7 @@ export UV_PROJECT_ENVIRONMENT := .venv-$(PY)
 endif
 PY_PATHS := src tests scripts examples .claude/hooks
 
-.PHONY: install fmt lint type test coverage check lowest smoke fuzz bench backends docs docs-serve notebook release
+.PHONY: install fmt lint type type-mypy type-pyright type-zuban type-pyrefly test coverage check lowest smoke fuzz bench backends docs docs-serve notebook release
 
 install:
 	$(UV) sync
@@ -38,8 +38,20 @@ lint:
 	# makes unclassifiable imports fail too
 	$(RUN) cleanporter $(PY_PATHS)
 
-type:
-	$(RUN) mypy
+type: type-mypy type-pyright type-zuban type-pyrefly
+
+type-mypy:
+	$(RUN) --frozen mypy
+
+# Query uv's interpreter so PY= and UV_PROJECT_ENVIRONMENT work too.
+type-pyright:
+	$(RUN) --frozen pyright --pythonpath "$$($(RUN) --frozen python -c 'import sys; print(sys.executable)')"
+
+type-zuban:
+	$(RUN) --frozen zuban mypy
+
+type-pyrefly:
+	$(RUN) --frozen pyrefly check --python-interpreter-path "$$($(RUN) --frozen python -c 'import sys; print(sys.executable)')"
 
 test:
 	$(RUN) pytest

docs/development.md

@@ -22,9 +22,11 @@ dependency at its floor (`--resolution lowest-direct`) on Python 3.12.
 ## Commands
 
 ```bash
-make check        # cleanporter + ruff + mypy (strict) + pytest + coverage (gate: 95%)
+make check        # cleanporter + ruff + four type checkers + pytest + coverage (95%)
 make check PY=3.13    # any target on another Python (env: .venv-3.13)
 make lint type test   # the same gates without coverage
+make type         # mypy + pyright + zuban + pyrefly
+make type-mypy    # individual checker (also: type-pyright, type-zuban, type-pyrefly)
 make lowest       # the test suite at the minimum dependency versions
 make fmt          # cleanporter --fix, then ruff format / check --fix / format
 make smoke        # run scripts/smoke.py, examples/demo.py, the notebooks
@@ -47,6 +49,45 @@ plain `pytest` skips nothing but keeps benchmarks disabled
 same check with `--fix`: cleanporter rewrites first, so ruff formats the
 already-wrapped lines (its module prefixes push lines past 100 columns).
 
+## Type checking
+
+All four checkers are blocking gates in `make type`, `make check`, git hooks,
+CI, and releases. Make owns their commands; `pyproject.toml` owns their scopes
+and settings. The recipes use `uv run --frozen`, so they run the locked versions
+without changing the lock. Hooks invoke the same Make targets and check the
+whole configured scope, even when only one file is staged. Python/stub files,
+the Makefile, hook configuration, and dependency/typing configuration trigger them.
+
+| Checker | Scope | Mode |
+|---|---|---|
+| mypy | `src` | strict |
+| Pyright | `src` | standard |
+| Zuban | `src` and maintained tests | strict, with test annotation requirements relaxed |
+| Pyrefly | `src` and maintained tests | default diagnostics, including unannotated bodies |
+
+All target Python 3.12. Zuban has its own configuration rather than inheriting
+mypy's source-only scope. The wider pair resolves `support` and `corpus` from
+the `tests` import root, like pytest. They exclude compiler input data
+(`tests/corpus.py`), generated golden cases, and the optional JAX/PyTorch
+compilation test module; the golden test harness and ordinary backend tests
+remain checked. JAX/PyTorch compilation tests run in the dedicated CI job.
+
+Pyrefly disables default hidden-path exclusions so checkouts inside `.paseo`
+or another hidden worktree directory still check tests. Run it in project
+mode: passing positional paths bypasses configured exclusions. Read checker
+output after changing configuration; some tools only warn about unknown keys.
+
+Pyright and Pyrefly receive uv's active interpreter, so `make type PY=3.13`
+and custom `UV_PROJECT_ENVIRONMENT` values resolve the correct environment.
+Pyright's PyPI wrapper downloads Node on its first run, which can take longer.
+
+Fix diagnostics with accurate annotations and narrowing. The compiler attaches
+metadata to functions dynamically: tests use `support.with_metadata` to describe
+that boundary without changing the public API. `support.binding` checks IR
+statement shapes before accessing binding-only fields. Deliberate invalid-input
+tests may need a narrowly documented typing boundary; do not hide ordinary code
+behind blanket suppressions.
+
 ## Testing layers
 
 | Layer | Where |
@@ -101,7 +142,8 @@ the NumPy-style docstring subject).
 
 ## CI
 
-- **check** — the full `make check` gate on Python 3.12–3.14
+- **static** — `make lint type` once on Python 3.12, including all four checkers.
+- **check** — `make coverage` on Python 3.12–3.14
   (Hypothesis derandomized), then `make smoke`.
 - **lowest** — `make lowest`: the suite at the declared dependency floors.
 - **fuzz** — the five fixed fuzzer seeds.

pyproject.toml

@@ -41,6 +41,9 @@ dev = [
     "hypothesis>=6.100",
     "ruff>=0.15",
     "mypy>=1.20",
+    "pyright>=1.1.411",
+    "zuban>=0.9.2",
+    "pyrefly>=1.2",
     "coverage[toml]>=7.6",
     "numpy>=2.0",
     "array-api-strict>=2.3",
@@ -163,6 +166,61 @@ files = ["src"]
 module = "array_api_compat"
 ignore_missing_imports = true
 
+# All four checkers are gates; Make owns commands and these sections own scope.
+[tool.pyright]
+include = ["src"]
+pythonVersion = "3.12"
+typeCheckingMode = "standard"
+
+[tool.zuban]
+files = ["src", "tests"]
+mypy_path = ["src", "tests"]
+python_version = "3.12"
+strict = true
+check_untyped_defs = true
+# Compiler input data and optional heavyweight integration tests are not
+# maintained test harness code. Zuban matches names during recursive discovery,
+# so patterns must also match the basename, not just a repository-relative path.
+exclude = [
+    '(^|/)cases(/|$)',
+    '(^|/)corpus\.py$',
+    '(^|/)test_jax_torch_compilation\.py$',
+]
+
+[[tool.zuban.overrides]]
+module = [
+    "tests.*", "unit.*", "behavior.*", "features.*", "backends.*",
+    "differential.*", "golden.*", "perf.*", "test_*", "support", "conftest",
+]
+disallow_untyped_defs = false
+disallow_incomplete_defs = false
+disallow_untyped_calls = false
+disallow_untyped_decorators = false
+
+[[tool.zuban.overrides]]
+# Excluding discovery alone does not prevent imports from following corpus data.
+module = ["corpus", "tests.corpus"]
+follow_imports = "skip"
+
+[[tool.zuban.overrides]]
+module = "array_api_compat"
+ignore_missing_imports = true
+
+[tool.pyrefly]
+project-includes = ["src", "tests"]
+project-excludes = [
+    "tests/golden/cases/**",
+    "tests/corpus.py",
+    "tests/backends/test_jax_torch_compilation.py",
+    "**/__pycache__/**",
+]
+search-path = ["src", "tests"]
+python-version = "3.12"
+check-unannotated-defs = true
+# This checkout can live under .paseo: default hidden-path exclusions skip tests.
+disable-project-excludes-heuristics = true
+replace-imports-with-any = ["corpus"]
+
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 pythonpath = ["tests"]

src/array_vectorize/codegen/emitter.py

@@ -208,7 +208,7 @@ def _gen_stmt(stmt: ir.Stmt, ns: str) -> ast.stmt:
                 value=_gen_expr(stmt.expr, ns),
             )
         case ir.Loop():
-            body = [_gen_stmt(inner, ns) for inner in stmt.body]
+            body: list[ast.stmt] = [_gen_stmt(inner, ns) for inner in stmt.body]
             if not body:
                 # (15) a retained loop with a fully dead body still needs a statement
                 body = [ast.Pass()]

src/array_vectorize/fallback.py

@@ -40,7 +40,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                 "Array API array argument"
             )
         arrays = [args[i] for i in positions] + [kwargs[k] for k in kw_positions]
-        xp = namespace if namespace is not None else array_api_compat.array_namespace(*arrays)
+        # The namespace is backend-dependent, including dynamically re-exported APIs.
+        xp: Any = namespace if namespace is not None else array_api_compat.array_namespace(*arrays)
         broadcast = xp.broadcast_arrays(*arrays)
         shape = broadcast[0].shape
         flat = [xp.reshape(a, (-1,)) for a in broadcast]
@@ -69,7 +70,8 @@ def with_namespace(xp: Any) -> Callable[..., Any]:
     # it is often unavailable here, which is a common fallback reason).
     # A pathological non-str __doc__ is left exactly as wraps left it.
     if isinstance(func.__doc__, str):
-        wrapper.__doc__ = docstring.prefixed_summary(func.__doc__)
+        # Zuban infers the nested wrapper's doc as None despite functools.wraps.
+        setattr(wrapper, "__doc__", docstring.prefixed_summary(func.__doc__))  # noqa: B010
     wrapper.source = (  # type: ignore[attr-defined]
         f"# fallback element-loop wrapper around {func.__name__!r}\n"
         f"# (source-to-source vectorization failed: {reason})"

src/array_vectorize/frontend/lambda_id.py

@@ -73,7 +73,8 @@ def _defaults_match(node: ast.Lambda, target: object) -> bool:
         if d is not None
     )
     fn_defaults = tuple(getattr(target, "__defaults__", None) or ())
-    fn_kw_defaults = tuple(v for k, v in (getattr(target, "__kwdefaults__", None) or {}).items())
+    kw_defaults: dict[str, Any] = getattr(target, "__kwdefaults__", None) or {}
+    fn_kw_defaults = tuple(kw_defaults.values())
     same = _const_key(ast_defaults) == _const_key(fn_defaults) and _const_key(
         ast_kw_defaults
     ) == _const_key(fn_kw_defaults)

src/array_vectorize/fuzz.py

@@ -24,7 +24,7 @@
 # load-bearing re-export: tests/unit/test_fuzz.py monkeypatches
 # array_vectorize.fuzz.vectorize, so the name must stay bound here
 from array_vectorize import errors
-from array_vectorize import vectorize  # cleanporter: ignore[CP003] load-bearing
+from array_vectorize import vectorize  # cleanporter: ignore[CP001] load-bearing re-export
 
 _UNARY_MATH = [
     f"math.{name}"

src/array_vectorize/lower/kinds.py

@@ -186,7 +186,7 @@ def record_carried_kind(self, name: str, kind: ir.Kind | None) -> None:
 
     def snapshot_facts(self) -> FactsSnapshot:
         """Capture kind + literal facts ONLY — never flags."""
-        kinds = {
+        kinds: dict[str, ir.Kind | None] = {
             name: info.kind for name, info in self._names.items() if name in self._kind_established
         }
         literals = {

src/array_vectorize/lower/loops.py

@@ -142,8 +142,10 @@ def union_labels(
             labels: dict[str, ir.Kind | None] = {}
             mixed: list[str] = []
             for loop_carried_name in carried.values():
-                kinds = {phi_kinds.get(loop_carried_name)} | frame.get(loop_carried_name, set())
-                unique = {k for k in kinds if k is not None}
+                kinds = set[ir.Kind | None]((phi_kinds.get(loop_carried_name),)) | frame.get(
+                    loop_carried_name, set()
+                )
+                unique: set[ir.Kind] = {k for k in kinds if k is not None}
                 if (None in kinds and unique) or len(unique) > 1:
                     # mixed (or unknown-mixed) kinds across iterations:
                     # label 'bool' so exact-arithmetic intify kicks in at

src/array_vectorize/runtime/arith.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 import enum
 import operator
 from typing import Any
@@ -46,7 +47,7 @@ class ArithOp(enum.IntEnum):
 #: (generated code passes the id as an int literal)
 _ARITH_OPS = tuple(m.name.lower() for m in ArithOp)
 
-_ARITH_FNS = {
+_ARITH_FNS: dict[str, Callable[[Any, Any], Any]] = {
     "add": operator.add,
     "sub": operator.sub,
     "mul": operator.mul,
@@ -197,12 +198,24 @@ def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any:
                 r = left % d
                 mag = q + xp.astype(r > 0, left.dtype)  # ceil(a / |d|)
                 return -xp.astype(mag, xp.int64)
-            if op == "mod" and isinstance(left, int) and left < 0 and _is_u64(right):
+            if (
+                op == "mod"
+                and isinstance(left, int)
+                and left < 0
+                and right is not None
+                and _is_u64(right)
+            ):
                 # negative-literal % array: (v - |d| mod v) mod v — always
                 # in [0, v), exactly representable in uint64
                 d = xp.asarray(-left, dtype=right.dtype)
                 return (right - (d % right)) % right
-            if op == "floordiv" and isinstance(left, int) and left < 0 and _is_u64(right):
+            if (
+                op == "floordiv"
+                and isinstance(left, int)
+                and left < 0
+                and right is not None
+                and _is_u64(right)
+            ):
                 # negative-literal // array: -ceil(|d| / v), exact in int64
                 d = xp.asarray(-left, dtype=right.dtype)
                 mag = (d // right) + xp.astype((d % right) > 0, right.dtype)
@@ -222,9 +235,8 @@ def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any:
                     return -xp.astype(left, xp.int64)
                 return -xp.astype(left, xp.float64)
         if op == "div":
-            return operator.truediv(
-                xp.asarray(left, dtype=xp.float64), xp.asarray(right, dtype=xp.float64)
-            )
+            divide: Callable[[Any, Any], Any] = operator.truediv
+            return divide(xp.asarray(left, dtype=xp.float64), xp.asarray(right, dtype=xp.float64))
     # generic: cast to the common dtype and compute
     classes = [dtype._describe_dtype(a) for a in args]
     has_float = any(f for f, _ in classes)

src/array_vectorize/verify.py

@@ -40,7 +40,8 @@ def verify_match(
             "verify= needs at least one Array API array in the example inputs"
         )
     arrays = [example_args[i] for i in positions]
-    xp = namespace if namespace is not None else array_api_compat.array_namespace(*arrays)
+    # The namespace is backend-dependent, including dynamically re-exported APIs.
+    xp: Any = namespace if namespace is not None else array_api_compat.array_namespace(*arrays)
     with warnings.catch_warnings():
         # out-of-domain lanes are expected (design D3); their warnings are noise
         # (including inf-inf in the comparison below)