From ac21cfcbd87894e89844153c32df6a40aedd745a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 11:38:46 +0200 Subject: [PATCH] ci: gate development with four python type checkers --- .github/workflows/ci.yml | 18 ++++- .pre-commit-config.yaml | 29 +++++++++ AGENTS.md | 12 +++- Makefile | 18 ++++- docs/development.md | 46 ++++++++++++- pyproject.toml | 58 +++++++++++++++++ src/array_vectorize/codegen/emitter.py | 2 +- src/array_vectorize/fallback.py | 6 +- src/array_vectorize/frontend/lambda_id.py | 3 +- src/array_vectorize/fuzz.py | 2 +- src/array_vectorize/lower/kinds.py | 2 +- src/array_vectorize/lower/loops.py | 6 +- src/array_vectorize/runtime/arith.py | 24 +++++-- src/array_vectorize/verify.py | 3 +- tests/behavior/test_helpers.py | 6 +- tests/behavior/test_inspectability.py | 13 ++-- tests/behavior/test_lambdas.py | 2 +- tests/behavior/test_loops.py | 12 ++-- tests/behavior/test_memoization.py | 15 +++-- tests/behavior/test_name_allocation.py | 2 +- tests/differential/test_scalar_equivalence.py | 14 +++- tests/features/test_fallback.py | 4 +- tests/features/test_namespace_pin.py | 65 +++++++++++-------- tests/features/test_protect_domains.py | 10 +-- tests/golden/test_golden.py | 5 +- tests/support.py | 37 ++++++++++- tests/unit/test_builtins.py | 3 +- tests/unit/test_extract.py | 7 +- tests/unit/test_fuzz.py | 3 +- tests/unit/test_lower.py | 22 +++---- tests/unit/test_optimize.py | 27 ++++---- tests/unit/test_validate.py | 6 +- uv.lock | 57 ++++++++++++++++ 33 files changed, 421 insertions(+), 118 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 99c4fff..e541c0e 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a577b82..000846e 100644 --- a/.pre-commit-config.yaml +++ b/.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: diff --git a/AGENTS.md b/AGENTS.md index efadcaf..dfe6d32 100644 --- a/AGENTS.md +++ b/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/`), diff --git a/Makefile b/Makefile index 2e5d073..00ff4e8 100644 --- a/Makefile +++ b/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 diff --git a/docs/development.md b/docs/development.md index 7222a01..4c091e9 100644 --- a/docs/development.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 47d4644..85ebea3 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/array_vectorize/codegen/emitter.py b/src/array_vectorize/codegen/emitter.py index 45c8961..e4dcbd2 100644 --- a/src/array_vectorize/codegen/emitter.py +++ b/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()] diff --git a/src/array_vectorize/fallback.py b/src/array_vectorize/fallback.py index aab7b7e..7744d9c 100644 --- a/src/array_vectorize/fallback.py +++ b/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})" diff --git a/src/array_vectorize/frontend/lambda_id.py b/src/array_vectorize/frontend/lambda_id.py index 3c77f9e..8ebcd40 100644 --- a/src/array_vectorize/frontend/lambda_id.py +++ b/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) diff --git a/src/array_vectorize/fuzz.py b/src/array_vectorize/fuzz.py index 5be2802..35ba905 100644 --- a/src/array_vectorize/fuzz.py +++ b/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}" diff --git a/src/array_vectorize/lower/kinds.py b/src/array_vectorize/lower/kinds.py index 7214a69..cbacb94 100644 --- a/src/array_vectorize/lower/kinds.py +++ b/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 = { diff --git a/src/array_vectorize/lower/loops.py b/src/array_vectorize/lower/loops.py index 9efabd2..64632fe 100644 --- a/src/array_vectorize/lower/loops.py +++ b/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 diff --git a/src/array_vectorize/runtime/arith.py b/src/array_vectorize/runtime/arith.py index b6dcadf..e89bb51 100644 --- a/src/array_vectorize/runtime/arith.py +++ b/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) diff --git a/src/array_vectorize/verify.py b/src/array_vectorize/verify.py index 08e1e64..a4cee27 100644 --- a/src/array_vectorize/verify.py +++ b/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) diff --git a/tests/behavior/test_helpers.py b/tests/behavior/test_helpers.py index 4c36355..8e3f295 100644 --- a/tests/behavior/test_helpers.py +++ b/tests/behavior/test_helpers.py @@ -27,7 +27,7 @@ def test_helper_call_composition() -> None: xs = np.asarray([3.0, 4.0]) got = vec(xs) assert np.allclose(got, np.sqrt(2 * xs**2)) - assert "square_vec(" in vec.source + assert "square_vec(" in support.with_metadata(vec).source # the helper itself is vectorized and memoized from array_vectorize import api @@ -66,9 +66,9 @@ def test_helper_with_loop_calling_helper() -> None: # the loop lives in the helper's generated source, not the caller's from array_vectorize import api - helper_src = api._HELPER_CACHE[(mod.addmul, False, None)].source + helper_src = support.with_metadata(api._HELPER_CACHE[mod.addmul, False, None]).source assert "for i in range(0, 3):" in helper_src - assert "addmul_vec" in vec.source + assert "addmul_vec" in support.with_metadata(vec).source def test_same_named_helper_does_not_shadow_caller() -> None: diff --git a/tests/behavior/test_inspectability.py b/tests/behavior/test_inspectability.py index d4c1c75..7bc98a7 100644 --- a/tests/behavior/test_inspectability.py +++ b/tests/behavior/test_inspectability.py @@ -8,6 +8,7 @@ import traceback import pytest +import support import array_vectorize @@ -30,8 +31,8 @@ def sigmoid_like(x: float, alpha: float = 1.0) -> float: def test_source_attribute() -> None: vec = array_vectorize.vectorize(relu_fn) - assert isinstance(vec.source, str) - assert "xp.where" in vec.source + assert isinstance(support.with_metadata(vec).source, str) + assert "xp.where" in support.with_metadata(vec).source def test_getsource_returns_generated_not_original() -> None: @@ -54,7 +55,9 @@ def test_linecache_survives_checkcache() -> None: filename = vec.__code__.co_filename assert filename == "" linecache.checkcache(filename) - assert linecache.cache[filename][2] # still registered + entry = linecache.cache[filename] + assert len(entry) == 4 # populated entry, rather than the one-item lazy entry + assert entry[2] # still registered assert "array_namespace" in "".join(linecache.getlines(filename)) @@ -74,7 +77,7 @@ def test_wrapped_points_to_original() -> None: # __wrapped__ is deliberately not set (it would hide the generated source # from inspect.getsourcelines, which unwraps unconditionally); the # original is reachable via the marker and drives the signature. - assert vec._vectorized_original is relu_fn + assert support.with_metadata(vec)._vectorized_original is relu_fn assert inspect.signature(vec) == inspect.signature(relu_fn) @@ -109,7 +112,7 @@ def with_defaults(x: float, scale: float = 2.0, *, bias: float = 0.5) -> float: def test_get_source_helper() -> None: vec = array_vectorize.vectorize(relu_fn) - assert array_vectorize.get_source(vec) == vec.source + assert array_vectorize.get_source(vec) == support.with_metadata(vec).source with pytest.raises(array_vectorize.VectorizationError, match="not a vectorized function"): array_vectorize.get_source(relu_fn) diff --git a/tests/behavior/test_lambdas.py b/tests/behavior/test_lambdas.py index 5676d70..bace100 100644 --- a/tests/behavior/test_lambdas.py +++ b/tests/behavior/test_lambdas.py @@ -14,7 +14,7 @@ def test_lambda_vectorize() -> None: vec = array_vectorize.vectorize(lambda x: x * 2.0) got = vec(np.asarray([1.0, 2.0])) assert np.allclose(got, [2.0, 4.0]) - assert "def lambda_vec" in vec.source + assert "def lambda_vec" in support.with_metadata(vec).source # ------------------------------------------------------- lambda identification diff --git a/tests/behavior/test_loops.py b/tests/behavior/test_loops.py index 5d43c52..2d32f50 100644 --- a/tests/behavior/test_loops.py +++ b/tests/behavior/test_loops.py @@ -19,7 +19,7 @@ def test_loop_with_start_stop_step() -> None: ) got = vec(np.asarray([1.0])) assert np.allclose(got, [(1 + 3) * 1.0]) - assert "for i in range(1, 5, 2):" in vec.source + assert "for i in range(1, 5, 2):" in support.with_metadata(vec).source def test_loop_negative_step() -> None: @@ -29,7 +29,7 @@ def test_loop_negative_step() -> None: ) ) assert np.allclose(vec(np.asarray([1.0])), [6.0]) - assert "for i in range(3, 0, -1):" in vec.source + assert "for i in range(3, 0, -1):" in support.with_metadata(vec).source def test_range_two_args() -> None: @@ -52,7 +52,7 @@ def test_loop_bound_expression_folds() -> None: vec = array_vectorize.vectorize( support.make_fn(" s = 0.0\n for i in range(1 + 2):\n s = s + x\n return s") ) - assert "for i in range(0, 3):" in vec.source + assert "for i in range(0, 3):" in support.with_metadata(vec).source # --------------------------------------------------------- loop-carried values @@ -64,8 +64,8 @@ def test_simple_accumulation() -> None: ) got = vec(np.asarray([1.0, 2.0])) assert np.allclose(got, [3.0, 6.0]) - assert "for i in range(0, 3):" in vec.source - assert "s_1 = s" in vec.source # loop-carried phi + assert "for i in range(0, 3):" in support.with_metadata(vec).source + assert "s_1 = s" in support.with_metadata(vec).source # loop-carried phi def test_loop_index_used() -> None: @@ -179,7 +179,7 @@ def test_loop_dce_drops_dead_body_bindings() -> None: ) vec = array_vectorize.vectorize(support.make_fn(body)) # the original source in the docstring mentions 'dead' once; the code must not - assert vec.source.count("dead") == 1 + assert support.with_metadata(vec).source.count("dead") == 1 def test_loop_dce_keeps_cross_iteration_dependencies() -> None: diff --git a/tests/behavior/test_memoization.py b/tests/behavior/test_memoization.py index a4b1cbb..bad0520 100644 --- a/tests/behavior/test_memoization.py +++ b/tests/behavior/test_memoization.py @@ -4,6 +4,7 @@ import numpy as np import pytest +import support import array_vectorize @@ -34,12 +35,12 @@ def helper_outer_fn(x: float) -> float: def test_backref_on_canonical_vectorization() -> None: vec = array_vectorize.vectorize(sigmoid_like) - assert sigmoid_like.__array_vectorized__ is vec + assert support.backref(sigmoid_like) is vec def test_vectorize_memoized_per_original() -> None: assert array_vectorize.vectorize(sigmoid_like) is array_vectorize.vectorize(sigmoid_like) - assert array_vectorize.vectorize(sigmoid_like) is sigmoid_like.__array_vectorized__ + assert array_vectorize.vectorize(sigmoid_like) is support.backref(sigmoid_like) def test_revectorization_returns_same_object() -> None: @@ -56,9 +57,9 @@ def test_verify_runs_on_cache_hit() -> None: def test_backref_not_overwritten_by_variants() -> None: vec = array_vectorize.vectorize(sigmoid_like) array_vectorize.vectorize(sigmoid_like, protect_domains=True) - assert sigmoid_like.__array_vectorized__ is vec + assert support.backref(sigmoid_like) is vec array_vectorize.vectorize(sigmoid_like, namespace=np) - assert sigmoid_like.__array_vectorized__ is vec + assert support.backref(sigmoid_like) is vec def test_backref_on_fallback_wrapper() -> None: @@ -69,7 +70,7 @@ def bad(x: float) -> float: with pytest.warns(UserWarning, match="falling back"): vec = array_vectorize.vectorize(bad, fallback=True) - assert bad.__array_vectorized__ is vec + assert support.backref(bad) is vec def test_helpers_do_not_get_backref() -> None: @@ -77,7 +78,7 @@ def test_helpers_do_not_get_backref() -> None: # must not grow an __array_vectorized__ marker array_vectorize.vectorize(helper_outer_fn) assert not hasattr(helper_inner_fn, "__array_vectorized__") - assert helper_outer_fn.__array_vectorized__ is not None + assert support.backref(helper_outer_fn) is not None def test_idempotent_revectorization() -> None: @@ -85,7 +86,7 @@ def test_idempotent_revectorization() -> None: vec2 = array_vectorize.vectorize(vec1) x = np.asarray([-1.0, 2.0]) assert np.allclose(vec1(x), vec2(x)) - assert "xp.where" in vec2.source + assert "xp.where" in support.with_metadata(vec2).source # canonical calls are memoized per scalar original: re-vectorizing # returns the very same object assert vec2 is vec1 diff --git a/tests/behavior/test_name_allocation.py b/tests/behavior/test_name_allocation.py index 5526fe7..e07db9f 100644 --- a/tests/behavior/test_name_allocation.py +++ b/tests/behavior/test_name_allocation.py @@ -23,7 +23,7 @@ def test_reserved_xp_param_keyword_call() -> None: mod = support.make_module("import numpy as np\n\n\ndef subject(xp):\n return xp + 1\n") vec = array_vectorize.vectorize(mod.subject) assert np.allclose(vec(xp=np.asarray([1.0])), [2.0]) - assert "def subject_vec(xp):" in vec.source + assert "def subject_vec(xp):" in support.with_metadata(vec).source def test_namespace_var_rebinding_no_collision() -> None: diff --git a/tests/differential/test_scalar_equivalence.py b/tests/differential/test_scalar_equivalence.py index 901e04e..011fac1 100644 --- a/tests/differential/test_scalar_equivalence.py +++ b/tests/differential/test_scalar_equivalence.py @@ -2,7 +2,8 @@ from __future__ import annotations -from typing import Any +from collections.abc import Callable +from typing import Any, cast import corpus import hypothesis @@ -30,7 +31,7 @@ def arrays( def check(fn: Any, xs: np.ndarray, *rest: np.ndarray) -> None: vec = array_vectorize.vectorize(fn) got = vec(xs, *rest) - pairs = list(zip(xs, *rest, strict=True)) if rest else [(x,) for x in xs] + pairs: list[tuple[Any, ...]] = list(zip(xs, *rest, strict=True)) if rest else [(x,) for x in xs] expected = np.asarray([fn(*pair) for pair in pairs]) assert np.allclose(got, expected, equal_nan=True, rtol=1e-12), (got, expected) @@ -67,11 +68,18 @@ def test_ternary_diff(xs: np.ndarray) -> None: check(CORPUS.ternary, xs) +def _nonsingular(pair: tuple[float, float]) -> bool: + return pair[0] * pair[0] + pair[1] * pair[1] > 1e-6 + + #: lanes where arith_ops' denominator x*x + y*y is not ~0; filtering per #: lane (not assume() over all 12) keeps Hypothesis from discarding most #: examples, since it draws 0.0 often +small_pairs: st.SearchStrategy[tuple[float, float]] = st.tuples(small, small) nonsingular_lanes = st.lists( - st.tuples(small, small).filter(lambda p: p[0] * p[0] + p[1] * p[1] > 1e-6), + # Hypothesis 6.100's conditional TypeVar definition confuses Pyrefly. + # Adapt only the callback boundary; the predicate keeps its tuple signature. + small_pairs.filter(cast(Callable[[Any], bool], _nonsingular)), min_size=12, max_size=12, ) diff --git a/tests/features/test_fallback.py b/tests/features/test_fallback.py index 3f14128..855b19e 100644 --- a/tests/features/test_fallback.py +++ b/tests/features/test_fallback.py @@ -58,7 +58,7 @@ def bad(x: float, scale: float = 2.0) -> float: with pytest.warns(UserWarning, match="falling back"): vec = array_vectorize.vectorize(bad, fallback=True) assert list(inspect.signature(vec).parameters) == ["x", "scale"] - assert "fallback" in vec.source + assert "fallback" in support.with_metadata(vec).source def test_fallback_non_str_docstring_left_verbatim() -> None: @@ -73,7 +73,7 @@ def bad(x: float) -> float: bad.__doc__ = 42 # type: ignore[assignment] with pytest.warns(UserWarning, match="falling back"): vec = array_vectorize.vectorize(bad, fallback=True) - assert vec.__doc__ == 42 + assert getattr(vec, "__doc__") == 42 # noqa: B009 - deliberately non-string metadata assert np.allclose(vec(np.asarray([5.0, -2.0])), [0.0, -2.0]) diff --git a/tests/features/test_namespace_pin.py b/tests/features/test_namespace_pin.py index 9c94545..44bfb1f 100644 --- a/tests/features/test_namespace_pin.py +++ b/tests/features/test_namespace_pin.py @@ -28,11 +28,12 @@ def test_pinned_source_shape() -> None: vec = array_vectorize.vectorize(corpus.add, namespace=np) - assert "def add_vec(x, y, *, _namespace=None):" in vec.source - assert "xp = _namespace" in vec.source + assert "def add_vec(x, y, *, _namespace=None):" in support.with_metadata(vec).source + assert "xp = _namespace" in support.with_metadata(vec).source # pinned generated source has zero imports and no namespace extraction - assert "import" not in vec.source - assert "array_namespace" not in vec.source + assert "import" not in support.with_metadata(vec).source + assert "array_namespace" not in support.with_metadata(vec).source + assert vec.__kwdefaults__ is not None assert vec.__kwdefaults__["_namespace"] is np @@ -41,15 +42,16 @@ def test_unpinned_source_unchanged() -> None: pinned = array_vectorize.vectorize(corpus.add, namespace=np) # byte-identical to the pre-pin generated source (the golden snapshot) golden = (pathlib.Path(__file__).parent.parent / "golden" / "cases" / "add.py").read_text() - assert unpinned.source == golden - assert pinned.source != unpinned.source - assert "xp = _namespace" not in unpinned.source - assert "_namespace=None" not in unpinned.source + assert support.with_metadata(unpinned).source == golden + assert support.with_metadata(pinned).source != support.with_metadata(unpinned).source + assert "xp = _namespace" not in support.with_metadata(unpinned).source + assert "_namespace=None" not in support.with_metadata(unpinned).source def test_pinned_closure_array_still_hidden_param() -> None: vec = array_vectorize.vectorize(corpus.closure_array, namespace=np) - assert "*, ARR=None, _namespace=None" in vec.source + assert "*, ARR=None, _namespace=None" in support.with_metadata(vec).source + assert vec.__kwdefaults__ is not None assert vec.__kwdefaults__["_namespace"] is np assert np.allclose(vec(np.asarray([0.0, 1.0, 2.0])), [1.0, 3.0, 5.0]) @@ -107,12 +109,14 @@ def test_all_scalar_call_without_pin_raises() -> None: def test_with_namespace_equivalent_results() -> None: vf = array_vectorize.vectorize(corpus.psi) - assert np.allclose(vf.with_namespace(np)(X), vf(X)) + assert np.allclose(support.with_metadata(vf).with_namespace(np)(X), vf(X)) def test_with_namespace_identity_stable() -> None: vf = array_vectorize.vectorize(corpus.psi) - assert vf.with_namespace(np) is vf.with_namespace(np) + assert support.with_metadata(vf).with_namespace(np) is support.with_metadata(vf).with_namespace( + np + ) def test_with_namespace_memoized_across_compilations() -> None: @@ -121,15 +125,17 @@ def test_with_namespace_memoized_across_compilations() -> None: vf1 = array_vectorize.vectorize(corpus.psi, protect_domains=True) vf2 = array_vectorize.vectorize(corpus.psi, protect_domains=True) assert vf1 is not vf2 - assert vf1.with_namespace(np) is vf2.with_namespace(np) + assert support.with_metadata(vf1).with_namespace(np) is support.with_metadata( + vf2 + ).with_namespace(np) def test_with_namespace_does_not_mutate() -> None: vf = array_vectorize.vectorize(corpus.psi) - w = vf.with_namespace(np) + w = support.with_metadata(vf).with_namespace(np) assert w is not vf # the original stays unpinned: no hidden namespace parameter at all - assert "_namespace=None" not in vf.source + assert "_namespace=None" not in support.with_metadata(vf).source assert (vf.__kwdefaults__ or {}).get("_namespace") is None @@ -138,16 +144,20 @@ def test_with_namespace_chaining_replaces_pin() -> None: vf = array_vectorize.vectorize(corpus.psi) a = xps.asarray([-2.0, 0.5, 3.0]) - chained = vf.with_namespace(np).with_namespace(xps) + chained = support.with_metadata(vf).with_namespace(np).with_namespace(xps) assert list(map(float, chained(a))) == list(map(float, vf(a))) def test_with_namespace_preserves_contract() -> None: vf = array_vectorize.vectorize(corpus.psi) - w = vf.with_namespace(np) + w = support.with_metadata(vf).with_namespace(np) assert w._vectorized_original is corpus.psi # the true scalar original assert inspect.signature(w) == inspect.signature(corpus.psi) - assert w.source == array_vectorize.vectorize(corpus.psi, namespace=np).source + assert ( + w.source + == support.with_metadata(array_vectorize.vectorize(corpus.psi, namespace=np)).source + ) + assert w.__kwdefaults__ is not None assert w.__kwdefaults__["_namespace"] is np assert "psi_vec" in inspect.getsource(w) # linecache registration intact @@ -163,7 +173,7 @@ def test_with_namespace_does_not_rerun_verify(monkeypatch: pytest.MonkeyPatch) - vec = array_vectorize.vectorize(fn, verify=(np.asarray([1.0]),)) assert calls # the initial decoration-time verify ran calls.clear() - vec.with_namespace(np) + support.with_metadata(vec).with_namespace(np) assert not calls # variants reuse verified body code, no re-verification @@ -199,8 +209,8 @@ def test_pinned_and_unpinned_helpers_are_distinct() -> None: unpinned = api._HELPER_CACHE[(mod.helper_inner, False, None)] pinned = api._HELPER_CACHE[(mod.helper_inner, False, np)] assert unpinned is not pinned - assert "xp = _namespace" in pinned.source - assert "array_namespace" in unpinned.source + assert "xp = _namespace" in support.with_metadata(pinned).source + assert "array_namespace" in support.with_metadata(unpinned).source # ---- name collisions @@ -210,8 +220,11 @@ def test_user_param_named_namespace_is_kept() -> None: fn = support.make_fn(" return x + _namespace", defaults="x, _namespace") vec = array_vectorize.vectorize(fn, namespace=np) # the user's parameter keeps its name; ours is mangled - assert "def subject_vec(x, _namespace, *, _namespace_1=None):" in vec.source - assert "xp = _namespace_1" in vec.source + assert ( + "def subject_vec(x, _namespace, *, _namespace_1=None):" in support.with_metadata(vec).source + ) + assert "xp = _namespace_1" in support.with_metadata(vec).source + assert vec.__kwdefaults__ is not None assert vec.__kwdefaults__["_namespace_1"] is np assert np.allclose(vec(np.asarray([1.0]), np.asarray([5.0])), [6.0]) # user keyword calls keep working @@ -221,8 +234,8 @@ def test_user_param_named_namespace_is_kept() -> None: def test_user_local_named_namespace_is_kept() -> None: fn = support.make_fn(" _namespace = x + 1\n return _namespace * 2", defaults="x") vec = array_vectorize.vectorize(fn, namespace=np) - assert "_namespace = x + 1" in vec.source - assert "xp = _namespace_1" in vec.source + assert "_namespace = x + 1" in support.with_metadata(vec).source + assert "xp = _namespace_1" in support.with_metadata(vec).source assert np.allclose(vec(np.asarray([1.0])), [4.0]) @@ -244,7 +257,7 @@ def test_fallback_with_namespace() -> None: with warnings.catch_warnings(): warnings.simplefilter("ignore") vec = array_vectorize.vectorize(fn, fallback=True) - pinned = vec.with_namespace(xps) + pinned = support.with_metadata(vec).with_namespace(xps) assert list(map(float, pinned(xps.asarray([3.0])))) == [0.0] @@ -268,7 +281,7 @@ def test_invalid_namespace_at_vectorize() -> None: def test_invalid_namespace_at_with_namespace() -> None: vf = array_vectorize.vectorize(support.make_fn(" return x + 1.0", defaults="x")) with pytest.raises(TypeError, match="not an Array API namespace"): - vf.with_namespace(object()) + support.with_metadata(vf).with_namespace(object()) # ---- verify= with a pin diff --git a/tests/features/test_protect_domains.py b/tests/features/test_protect_domains.py index 7577d0d..1cd9162 100644 --- a/tests/features/test_protect_domains.py +++ b/tests/features/test_protect_domains.py @@ -15,7 +15,7 @@ def test_protect_domains_clamps_dead_lanes() -> None: fn = support.make_fn(" if x > 0:\n return math.sqrt(x)\n return 0.0") vec = array_vectorize.vectorize(fn, protect_domains=True) - assert "xp.logical_not(x > 0)" in vec.source # dead-lane guard + assert "xp.logical_not(x > 0)" in support.with_metadata(vec).source # dead-lane guard xs = np.asarray([-1.0, 0.5, 4.0]) with np.errstate(invalid="ignore", divide="ignore"): got = vec(xs) @@ -53,8 +53,8 @@ def test_protect_domains_two_sided_clamp() -> None: fn = support.make_fn(" if x != 0:\n return math.atanh(x)\n return 0.0") vec = array_vectorize.vectorize(fn, protect_domains=True) # dead-lane where-clamps (preserve -0.0 and live values, unlike maximum) - assert "xp.logical_not" in vec.source - assert "xp.maximum" not in vec.source + assert "xp.logical_not" in support.with_metadata(vec).source + assert "xp.maximum" not in support.with_metadata(vec).source xs = np.asarray([0.5, 0.0]) assert np.allclose(vec(xs), [math.atanh(0.5), 0.0]) @@ -101,13 +101,13 @@ def test_protect_domains_preserves_subnormal_live_lanes() -> None: def test_protect_domains_off_by_default() -> None: fn = support.make_fn(" if x > 0:\n return math.sqrt(x)\n return 0.0") vec = array_vectorize.vectorize(fn) - assert "xp.maximum(x, 0.0)" not in vec.source + assert "xp.maximum(x, 0.0)" not in support.with_metadata(vec).source def test_protect_domains_top_level_calls_not_clamped() -> None: fn = support.make_fn(" return math.sqrt(abs(x))") vec = array_vectorize.vectorize(fn, protect_domains=True) - assert "maximum" not in vec.source + assert "maximum" not in support.with_metadata(vec).source def test_protect_domains_conditions_not_clamped() -> None: diff --git a/tests/golden/test_golden.py b/tests/golden/test_golden.py index b81dbaf..1848f5e 100644 --- a/tests/golden/test_golden.py +++ b/tests/golden/test_golden.py @@ -8,6 +8,7 @@ import corpus import numpy as np import pytest +import support import array_vectorize @@ -20,12 +21,12 @@ def _generated_source(name: str) -> str: fn = getattr(corpus, name) - return array_vectorize.vectorize(fn).source + return support.with_metadata(array_vectorize.vectorize(fn)).source def _generated_pinned_source(name: str) -> str: fn = getattr(corpus, name) - return array_vectorize.vectorize(fn, namespace=np).source + return support.with_metadata(array_vectorize.vectorize(fn, namespace=np)).source def _golden_path(name: str) -> pathlib.Path: diff --git a/tests/support.py b/tests/support.py index e37783b..413e278 100644 --- a/tests/support.py +++ b/tests/support.py @@ -15,17 +15,50 @@ import itertools import pathlib import tempfile -from typing import Any +from typing import Any, cast, Protocol import corpus import array_vectorize +from array_vectorize import ir _tmp = tempfile.TemporaryDirectory(prefix="vec_support_") TMPDIR = pathlib.Path(_tmp.name) _seq = itertools.count() +class Vectorized(Protocol): + """Metadata attached dynamically to compiled functions, for test assertions.""" + + source: str + _vectorized_original: Callable[..., Any] + __kwdefaults__: dict[str, Any] | None + + def __call__(self, *args: Any, **kwargs: Any) -> Any: ... + + def with_namespace(self, xp: Any) -> Vectorized: ... + + +def with_metadata(fn: Callable[..., Any]) -> Vectorized: + """Describe compiler-attached attributes without changing the public API. + + Tests asserting metadata use this boundary; ordinary calls still retain + their callable types, and the assertions check the attributes at runtime. + """ + return cast(Vectorized, fn) + + +def binding(stmt: ir.Stmt) -> ir.Binding: + """Assert a statement is a binding before checking its name or expression.""" + assert isinstance(stmt, ir.Binding) + return stmt + + +def backref(fn: Callable[..., Any]) -> Callable[..., Any]: + """Read the dynamically attached vectorization marker on a scalar function.""" + return cast(Callable[..., Any], getattr(fn, "__array_vectorized__")) # noqa: B009 + + def make_fn(body: str, extra: str = "", defaults: str = "x, y=2.0") -> Callable[..., Any]: """Write a snippet module and return its ``subject`` function. @@ -47,7 +80,7 @@ def make_fn(body: str, extra: str = "", defaults: str = "x, y=2.0") -> Callable[ assert spec.loader is not None mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) - return mod.subject + return cast(Callable[..., Any], mod.subject) def make_module(source: str) -> Any: diff --git a/tests/unit/test_builtins.py b/tests/unit/test_builtins.py index 8eb5910..37a183d 100644 --- a/tests/unit/test_builtins.py +++ b/tests/unit/test_builtins.py @@ -38,7 +38,8 @@ def test_round_decimals_kwarg_is_not_standard() -> None: # array-api-strict is the conformance reference; it must reject it. x = xps.asarray([1.23456]) try: - xps.round(x, decimals=2) + # Deliberately call an unsupported backend keyword to test its rejection. + getattr(xps, "round")(x, decimals=2) # noqa: B009 except TypeError: pass else: diff --git a/tests/unit/test_extract.py b/tests/unit/test_extract.py index cd88e7a..1be83af 100644 --- a/tests/unit/test_extract.py +++ b/tests/unit/test_extract.py @@ -3,6 +3,7 @@ from __future__ import annotations import math +from typing import Any import numpy as np import pytest @@ -182,7 +183,7 @@ def f(x: float) -> float: def test_closure_array_captured() -> None: - def f(x: float) -> float: + def f(x: float) -> Any: return x + ARR[0] * 0.0 + x * 0.0 + x * 1.0 # ARR referenced (subscript rejected later) info = extract.extract_function(f) @@ -198,7 +199,7 @@ def f(x: float) -> float: def test_unbound_name_rejected() -> None: - def f(x: float) -> float: + def f(x: float) -> Any: return x + unknown_name # type: ignore[name-defined] # noqa: F821 with pytest.raises(errors.VectorizationError, match="not resolvable"): @@ -228,7 +229,7 @@ def f(x: float, x_1: float) -> float: def test_np_float64_is_scalar() -> None: c = np.float64(1.5) - def f(x: float) -> float: + def f(x: float) -> np.float64: return x * c info = extract.extract_function(f) diff --git a/tests/unit/test_fuzz.py b/tests/unit/test_fuzz.py index 93da350..28a25b7 100644 --- a/tests/unit/test_fuzz.py +++ b/tests/unit/test_fuzz.py @@ -14,7 +14,8 @@ def test_fuzzer_detects_mismatch(monkeypatch) -> None: # type: ignore[no-untype # sabotage the oracle to prove the fuzzer fails loudly on mismatch import array_vectorize.fuzz as fz - original = fz.vectorize + # The imported binding is a monkeypatch seam, not an explicit module export. + original = getattr(fz, "vectorize") # noqa: B009 def lying_vectorize(fn, **kwargs): # type: ignore[no-untyped-def] vec = original(fn, **kwargs) diff --git a/tests/unit/test_lower.py b/tests/unit/test_lower.py index cb37e8f..5339926 100644 --- a/tests/unit/test_lower.py +++ b/tests/unit/test_lower.py @@ -4,7 +4,7 @@ from collections.abc import Callable import math -from typing import Any +from typing import Any, cast import numpy as np import pytest @@ -31,7 +31,7 @@ def make_fn(body: str, extra_globals: str = "") -> Callable[..., Any]: + body + "\n" ) - return mod.subject + return cast(Callable[..., Any], mod.subject) def lower(body: str) -> ir.Program: @@ -197,8 +197,8 @@ def test_boolop_numeric_chain_uses_temp() -> None: p = lower(" return x and y and x") assert len(p.bindings) == 1 temp = p.bindings[0] - assert temp.name == "v_1" - assert temp.expr == ir.Where( + assert support.binding(temp).name == "v_1" + assert support.binding(temp).expr == ir.Where( ir.Compare("ne", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("y"), ir.Ref("x") ) assert p.result == ir.Where( @@ -239,7 +239,7 @@ def test_augassign() -> None: def test_rebind_skips_user_names() -> None: p = lower(" x_1 = 0\n x = x + 1\n x = x + 1\n return x + x_1") - names = [b.name for b in p.bindings] + names = [support.binding(b).name for b in p.bindings] assert names[1:] == ["x_2", "x_3"] # x_1 taken by the user variable @@ -308,7 +308,7 @@ def test_calling_variable_rejected() -> None: def test_if_else_merge() -> None: p = lower(" if x > 0:\n r = x\n else:\n r = 0.0\n return r") - assert [b.name for b in p.bindings] == ["r_1", "r_2", "r"] + assert [support.binding(b).name for b in p.bindings] == ["r_1", "r_2", "r"] assert p.bindings[0] == ir.Binding("r_1", ir.Ref("x")) assert p.bindings[1] == ir.Binding("r_2", ir.Literal(0.0, "float")) assert p.bindings[2] == ir.Binding( @@ -320,15 +320,15 @@ def test_if_else_merge() -> None: def test_if_else_merge_with_prior_binding() -> None: p = lower(" r = 0.0\n if x > 0:\n r = x\n else:\n r = 1.0\n return r") - assert [b.name for b in p.bindings] == ["r", "r_1", "r_2", "r_3"] + assert [support.binding(b).name for b in p.bindings] == ["r", "r_1", "r_2", "r_3"] assert p.result == ir.Ref("r_3") def test_if_without_else_and_prior_binding() -> None: p = lower(" r = 0.0\n if x > 0:\n r = x\n return r") # then binds r_1; else keeps r; merge where(cond, r_1, r) - assert [b.name for b in p.bindings] == ["r", "r_1", "r_2"] - assert p.bindings[2].expr == ir.Where( + assert [support.binding(b).name for b in p.bindings] == ["r", "r_1", "r_2"] + assert support.binding(p.bindings[2]).expr == ir.Where( ir.Compare("gt", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("r_1"), ir.Ref("r") ) assert p.result == ir.Ref("r_2") @@ -344,7 +344,7 @@ def test_elif_chain() -> None: " r = 0.0\n" " return r" ) - names = [b.name for b in p.bindings] + names = [support.binding(b).name for b in p.bindings] assert names == ["r_1", "r_2", "r_3", "r", "r_4"] assert p.result == ir.Ref("r_4") @@ -500,7 +500,7 @@ def test_maybe_unbound_not_read_is_fine() -> None: p = lower(" if x > 0:\n z = 1.0\n return x") assert p.result == ir.Ref("x") # z_1 is dead; DCE (in _optimize) drops it later - assert [b.name for b in p.bindings] == ["z_1"] + assert [support.binding(b).name for b in p.bindings] == ["z_1"] def test_not_all_paths_return_rejected() -> None: diff --git a/tests/unit/test_optimize.py b/tests/unit/test_optimize.py index 87d16ce..0caa549 100644 --- a/tests/unit/test_optimize.py +++ b/tests/unit/test_optimize.py @@ -3,8 +3,10 @@ from __future__ import annotations import math +from typing import Literal import pytest +import support from array_vectorize import ir from array_vectorize import optimize as optimize_mod @@ -16,7 +18,7 @@ from array_vectorize.optimize import dce as dce_pass # cleanporter: ignore[CP002] name collision -def lit(v: float | int, kind: str = "auto") -> ir.Literal: +def lit(v: float | int, kind: ir.Kind | Literal["auto"] = "auto") -> ir.Literal: if kind == "auto": kind = "bool" if isinstance(v, bool) else "int" if isinstance(v, int) else "float" return ir.Literal(v, kind) @@ -57,7 +59,7 @@ def test_fold_pow_skipped() -> None: def test_fold_int64_overflow_skipped() -> None: - expr = ir.BinOp("mul", lit(2**62), lit(4)) + expr: ir.Node = ir.BinOp("mul", lit(2**62), lit(4)) assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr expr = ir.UnaryOp("neg", lit(-(2**63))) assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr @@ -125,7 +127,7 @@ def test_dce_keeps_chains() -> None: ), ir.Ref("b"), ) - assert [b.name for b in dce(p).bindings] == ["a", "b"] + assert [support.binding(b).name for b in dce(p).bindings] == ["a", "b"] def test_dce_single_reverse_pass_handles_late_uses() -> None: @@ -137,7 +139,7 @@ def test_dce_single_reverse_pass_handles_late_uses() -> None: ), ir.Ref("b"), ) - assert [b.name for b in dce(p).bindings] == ["a", "b"] + assert [support.binding(b).name for b in dce(p).bindings] == ["a", "b"] # ----------------------------------------------------------------------- CSE @@ -151,8 +153,8 @@ def test_cse_hoists_duplicated_call() -> None: ir.BinOp("mul", inner, ir.Ref("a")), ) out = cse(p, ir.SSAEnv(set())) - assert out.bindings[0].name == "t_1" - assert out.bindings[0].expr == inner + assert support.binding(out.bindings[0]).name == "t_1" + assert support.binding(out.bindings[0]).expr == inner assert out.bindings[1] == ir.Binding("a", ir.BinOp("add", ir.Ref("t_1"), ir.Ref("t_1"))) assert out.result == ir.BinOp("mul", ir.Ref("t_1"), ir.Ref("a")) @@ -179,7 +181,7 @@ def test_cse_nested_fixpoint() -> None: outer1 = ir.Where(ir.Compare("gt", inner, lit(0)), inner, lit(0.0, "float")) p = ir.Program(("x",), (), ir.BinOp("add", outer1, outer1)) out = cse(p, ir.SSAEnv(set())) - names = [b.name for b in out.bindings] + names = [support.binding(b).name for b in out.bindings] # inner temp created first, then (next pass) the rebuilt outer assert names == ["t_1", "t_2"] assert out.result == ir.BinOp("add", ir.Ref("t_2"), ir.Ref("t_2")) @@ -196,7 +198,7 @@ def test_cse_temp_inserted_before_first_user() -> None: ir.Ref("b"), ) out = cse(p, ir.SSAEnv(set())) - assert [b.name for b in out.bindings] == ["a", "t_1", "b"] + assert [support.binding(b).name for b in out.bindings] == ["a", "t_1", "b"] # ----------------------------------------------------------------- composite @@ -215,8 +217,8 @@ def test_optimize_pipeline() -> None: ) out = optimize_mod.optimize(p, {"x", "dead", "a"}) folded_sub = ir.BinOp("add", ir.Ref("x"), lit(6)) - assert out.bindings[0].name == "t_1" - assert out.bindings[0].expr == ir.Call("abs", (folded_sub,)) + assert support.binding(out.bindings[0]).name == "t_1" + assert support.binding(out.bindings[0]).expr == ir.Call("abs", (folded_sub,)) assert out.bindings[1] == ir.Binding("a", ir.Ref("t_1")) assert out.result == ir.BinOp("add", ir.Ref("t_1"), ir.Ref("a")) @@ -258,7 +260,7 @@ def test_fold_unary_ops_all() -> None: def test_fold_skips_bad_ops() -> None: # float bitwise -> TypeError at runtime -> no fold - expr = ir.BinOp("and", lit(1.5), lit(2.5)) + expr: ir.Node = ir.BinOp("and", lit(1.5), lit(2.5)) assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr # unary invert on float -> no fold expr = ir.UnaryOp("invert", lit(1.5)) @@ -274,7 +276,8 @@ def test_fold_skips_bad_ops() -> None: def test_fold_zero_times_infinity() -> None: # inf * 0 -> nan: exact IEEE fold out = optimize_mod.const_fold(ir.Program((), (), ir.BinOp("mul", lit(math.inf), lit(0.0)))) - assert math.isnan(out.result.value) # type: ignore[attr-defined] + assert isinstance(out.result, ir.Literal) + assert math.isnan(out.result.value) def test_dce_keeps_result_only_binding() -> None: diff --git a/tests/unit/test_validate.py b/tests/unit/test_validate.py index fcbabf4..15bbc70 100644 --- a/tests/unit/test_validate.py +++ b/tests/unit/test_validate.py @@ -5,7 +5,7 @@ import ast from collections.abc import Callable import sys -from typing import Any +from typing import Any, cast import pytest import support @@ -21,7 +21,7 @@ def make_fn(body: str, signature: str = "x", type_params: str = "") -> Callable[ mod = support.make_module( "import math\n\n\ndef subject" + type_params + "(" + signature + "):\n" + body + "\n" ) - return mod.subject + return cast(Callable[..., Any], mod.subject) def check(src: str) -> None: @@ -194,7 +194,7 @@ def subject(x: float) -> float: def test_all_violations_collected() -> None: - def subject(x: float) -> float: + def subject(x: Any) -> tuple[Any, Any]: y = x[0] while x > 0: x = x - 1 diff --git a/uv.lock b/uv.lock index 558cf58..991966c 100644 --- a/uv.lock +++ b/uv.lock @@ -68,10 +68,13 @@ dev = [ { name = "mypy" }, { name = "numpy" }, { name = "pre-commit" }, + { name = "pyrefly" }, + { name = "pyright" }, { name = "pytest" }, { name = "pytest-benchmark" }, { name = "pytest-timeout" }, { name = "ruff" }, + { name = "zuban" }, ] docs = [ { name = "zensical" }, @@ -98,10 +101,13 @@ dev = [ { name = "mypy", specifier = ">=1.20" }, { name = "numpy", specifier = ">=2.0" }, { name = "pre-commit", specifier = ">=3.5" }, + { name = "pyrefly", specifier = ">=1.2" }, + { name = "pyright", specifier = ">=1.1.411" }, { name = "pytest", specifier = ">=9.0" }, { name = "pytest-benchmark", specifier = ">=4.0" }, { name = "pytest-timeout", specifier = ">=2.3" }, { name = "ruff", specifier = ">=0.15" }, + { name = "zuban", specifier = ">=0.9.2" }, ] docs = [{ name = "zensical", specifier = ">=0.0.67" }] notebooks = [ @@ -1420,6 +1426,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a6/43/9f45ec4d14e596efc32c925a78104934790438b0c0628b70d741016734ad/pymdown_extensions-11.0.2-py3-none-any.whl", hash = "sha256:259910762019732caa1dfd76f3faa62c59f191d46573e80bcb1d13c0f675bbe5", size = 269929, upload-time = "2026-08-22T19:28:45.389Z" }, ] +[[package]] +name = "pyrefly" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4d/78/b6ec12acb586606618c71ebe1008af59e1ca50e667895f1946fa7b656619/pyrefly-1.3.2.tar.gz", hash = "sha256:30a68ff5429a1546ec08a2d913dfe442c9553a2940d274de18833d8eda939b5e", size = 6679591, upload-time = "2026-09-28T22:10:06.292Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/01/144a78839822abad55628881bcb45f8173e6d8b87976897c9105884b5a6b/pyrefly-1.3.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:626dbaee16ae1169a10a4ba3293583b5ab0320102f35efaad9e03457a2fc646c", size = 15088488, upload-time = "2026-09-28T22:09:40.743Z" }, + { url = "https://files.pythonhosted.org/packages/23/aa/599756bd42e6533ee9261ebd288f846ab453b00debaef2b4789893967809/pyrefly-1.3.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:2d5fc704dbf0e990dc4f15cb1ead64f1a325ae27383a50c76cfcaf0387f926bb", size = 14429996, upload-time = "2026-09-28T22:09:43.316Z" }, + { url = "https://files.pythonhosted.org/packages/90/0c/4b0678889136b1514ecd25cec11974d5bfdaeaaf9f5fe8a766ca45ec4ab5/pyrefly-1.3.2-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:89dbe8a33e735fa96023c33ce5dfd81af9b0bdd2ee771e8d69b2ebffc7e0fbc7", size = 14871810, upload-time = "2026-09-28T22:09:45.42Z" }, + { url = "https://files.pythonhosted.org/packages/a6/4e/a96236391c6bb0d8b50af85e0ea6c0c56206f739931ac3031f9e000325ee/pyrefly-1.3.2-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3ec0c31715d48bff279f548b85a93b5bc43d2c7bd1492a8cd703a9185cb089e5", size = 16113462, upload-time = "2026-09-28T22:09:47.619Z" }, + { url = "https://files.pythonhosted.org/packages/ad/c3/d4d0966aa7008ebd8481a6d05aa53202ad40c9597705eeee87ba6e0e2025/pyrefly-1.3.2-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3b3de468bed691e40dc6271284c8dcc05e649d7ed45280ef3d07332741efe34c", size = 16043341, upload-time = "2026-09-28T22:09:50.083Z" }, + { url = "https://files.pythonhosted.org/packages/4f/d1/52d9ea1b7048c7316d9ff7594c170b1c06d37547a75d140c8deea7b9ddfd/pyrefly-1.3.2-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cb94365fe5f854cb4fe82f600c4e2453160d328610b12981004851ee0b6e8aeb", size = 15437183, upload-time = "2026-09-28T22:09:52.56Z" }, + { url = "https://files.pythonhosted.org/packages/08/ec/24f83e4a397fa9d5c416bb4d378ee1bb18fe48bdd5698fe541636f5352f3/pyrefly-1.3.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:4f730d2e515ac25827fc68f343a7aa02b1072e0062df6a430346a4ffdfb5db2f", size = 14892976, upload-time = "2026-09-28T22:09:54.916Z" }, + { url = "https://files.pythonhosted.org/packages/29/c7/a473f8e9ec31ba0682ebe44f4e51167a22a4cbdc8bfd87365b9aca23be20/pyrefly-1.3.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:60f440e1626bd5cc4a41ee90cecba80796ed01ebd6c15d9f933efb8d07ffe2b8", size = 15469689, upload-time = "2026-09-28T22:09:57.541Z" }, + { url = "https://files.pythonhosted.org/packages/d0/df/779092365e8e469fc8e7349e1c26c35b8a73df9e6be80a00265fc26a2b2b/pyrefly-1.3.2-py3-none-win32.whl", hash = "sha256:b31684189de44a25a5bf00bdf3ad16be44a7b14654702d6018bce1043357202c", size = 14135550, upload-time = "2026-09-28T22:09:59.757Z" }, + { url = "https://files.pythonhosted.org/packages/9d/20/466333e7f2891f8c357b5c226559c39ba8dde4f90be09189ecc771a086f5/pyrefly-1.3.2-py3-none-win_amd64.whl", hash = "sha256:693960cdb247d1f76b9e09b3f2554dae20034f7d442ccd4e9b3e8962a3622363", size = 15067523, upload-time = "2026-09-28T22:10:02.099Z" }, + { url = "https://files.pythonhosted.org/packages/18/32/ee66df92c692c0686d82c9ecef3876ccb0386aa318bf507016f0121a8f78/pyrefly-1.3.2-py3-none-win_arm64.whl", hash = "sha256:f7424b21331670522980ff13e08bd143659879a4d92f4832282c7d7c41a9b5d6", size = 14348271, upload-time = "2026-09-28T22:10:04.253Z" }, +] + +[[package]] +name = "pyright" +version = "1.1.414" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nodeenv" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e1/1b/244c7b710031ada80f27e579ec20d28a2285dfc318fed0339866b1047f12/pyright-1.1.414.tar.gz", hash = "sha256:523c0a97c60da6333234955c277730c9cf4f5bd6d5399e7b7d2b0fc5d3599524", size = 4154638, upload-time = "2026-09-10T12:26:53.181Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/ba/18b6e682ead424ad24bcc134339ae5d1b931cd9ae260540592a058a91279/pyright-1.1.414-py3-none-any.whl", hash = "sha256:2a6b4b3298c9eec174c5ed83bd338de6eee82df2992f3e1930e6199d381be36f", size = 6225049, upload-time = "2026-09-10T12:26:51.427Z" }, +] + [[package]] name = "pytest" version = "9.1.1" @@ -2061,3 +2099,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6e/0d/84beca8c541326c4b0541f9f863f90a918cd3a57d02cd95701cfc6f2c8cd/zensical-0.0.68-cp311-abi3-win_amd64.whl", hash = "sha256:cc2cfdaa046fd56f0f38684c069078615564e3dab1688aba54ae71681b3a0270", size = 18716085, upload-time = "2026-10-05T15:41:34.14Z" }, { url = "https://files.pythonhosted.org/packages/ba/c8/ea991131b0fb72eaf46679f7b38c9c2285d584bfd5e3e138b868f509d44a/zensical-0.0.68-cp311-abi3-win_arm64.whl", hash = "sha256:f47a396ca7a7c60eefd5bb090f8974c94918937c7f7f41c316ecefd51eeec876", size = 18137272, upload-time = "2026-10-05T15:41:36.68Z" }, ] + +[[package]] +name = "zuban" +version = "0.10.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/e3/69ef2ec901c4a21b584c8f30e2957e6ee7706fff2b7f7b64554628c8e1da/zuban-0.10.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:59d7a698dd58bc742a6be187d000367d2dea541a2ee57531d10f4713520ee97c", size = 11352888, upload-time = "2026-09-18T11:11:30.344Z" }, + { url = "https://files.pythonhosted.org/packages/45/6b/726e27df1b5243fda6f2206147c6566a7e17e5ff64f8fec51c01e32d69c3/zuban-0.10.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c81631c3e6ba5e3455d3953588c89612d8d0995be28909824df44fa997af7e40", size = 11112654, upload-time = "2026-09-18T11:11:33.029Z" }, + { url = "https://files.pythonhosted.org/packages/99/36/ab87b7e4b3dcd5ca89ad324c9a5249b00892ed25a2c2c29ffdea1286650f/zuban-0.10.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bd4f6fcdd46d89f5fe5ca9191e05f3c1b6e28295b0bc532e18c3164565288380", size = 28554642, upload-time = "2026-09-18T11:11:35.821Z" }, + { url = "https://files.pythonhosted.org/packages/59/7a/82a17d92dc4652086356ed5642b4b81d285dd7f851b949ca967621d62e39/zuban-0.10.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:98fbd6e08e79d333cca612d4ed3f679a9cddbfc3b99f07693553769fe172035a", size = 28732309, upload-time = "2026-09-18T11:11:39.495Z" }, + { url = "https://files.pythonhosted.org/packages/bf/47/21116e950934fd040a10df5309978d0c2bb1084d2c2b04c668273422377a/zuban-0.10.0-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c494f530034b56b1ca9ef6d9851b793c858fa8697c6c01f683dc0ab203e290d1", size = 29838795, upload-time = "2026-09-18T11:11:43.292Z" }, + { url = "https://files.pythonhosted.org/packages/f6/4e/fa2148a3bac24788bda3a83c58167ae86e88b87a141911140953914d6447/zuban-0.10.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ae1fb50593f129a5df4dc0b06f25124345f8d527623545e1a79ffb451a78e131", size = 31846103, upload-time = "2026-09-18T11:11:47.362Z" }, + { url = "https://files.pythonhosted.org/packages/40/dd/6124f00f1e38c63ea548bc2a9dde5f547485e9119dbd1a320de5c54e0bdf/zuban-0.10.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:b08b5331728b818b87aae7a287f81f7d66bed56823a1084776ccdbc84067ffdb", size = 28725088, upload-time = "2026-09-18T11:11:50.857Z" }, + { url = "https://files.pythonhosted.org/packages/4a/87/5bef3b23e53f9cc525e5cd228c51d2e6f1fa47080f48b0d3ee0bb146af07/zuban-0.10.0-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:a66d60f6afd1deddef673dab407abf1014721e393fcc18fb1b5f2f73fa3f04d8", size = 29166127, upload-time = "2026-09-18T11:11:54.567Z" }, + { url = "https://files.pythonhosted.org/packages/00/5c/b173cbcfbf8261b27f6a70a54a92d3400a378a4143041073834c2c9c5cbc/zuban-0.10.0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:d83400da488b8a439358425ddfc800fa12859c03c8f7bcd8754e2cc487defea5", size = 29732954, upload-time = "2026-09-18T11:11:58.192Z" }, + { url = "https://files.pythonhosted.org/packages/f2/53/1b4d0a82aa728090458e531207a6a7634f186a566f5e0d5167ddc8f4355a/zuban-0.10.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:64002bf171baf49097b61521f7756a1c29264a3fcd2c98cf3e2b0b7862e74008", size = 29059999, upload-time = "2026-09-18T11:12:01.793Z" }, + { url = "https://files.pythonhosted.org/packages/ba/6d/2215ab53d7388a8d9de888af048010ef5b6421c54a50f7a38069e8e614ad/zuban-0.10.0-py3-none-win32.whl", hash = "sha256:9d39da51b2d679505011bdcf8ba5757b10dc3e24b941d572b6cde8217fa48539", size = 10002018, upload-time = "2026-09-18T11:12:05.125Z" }, + { url = "https://files.pythonhosted.org/packages/58/42/fb3c2c3a039a91608527a70beeddf1d183abdc8697bd124df8797676cefa/zuban-0.10.0-py3-none-win_amd64.whl", hash = "sha256:de928ef243a3547fa747321fb6aad4a07657827e72c23c7905dbeab4697f8337", size = 10713422, upload-time = "2026-09-18T11:12:07.69Z" }, +]