From 97b903d02af94beca276a51d78b98d6b1f3647c5 Mon Sep 17 00:00:00 2001 From: Jayson Steffens Date: Mon, 20 Apr 2026 16:20:11 +0200 Subject: [PATCH 1/3] chore: fix all 17 mypy --strict errors; promote to hard CI gate Resolved all outstanding mypy --strict errors so the warning-only check added in PR #43 can become a required gate. Changes: - Add snapvec/_fast.pyi stub so mypy sees the compiled Cython kernels' public API without resorting to 'type: ignore[import-not-found]'. Removes both the import errors and the unused-type-ignore warnings on the fallback branch in _pq.py / _ivfpq.py. - Annotate the five save-path closures (_write(f: 'ChecksumWriter') in _index.py, _pq.py, _residual.py, _ivfpq.py) and ChecksumWriter.__exit__ in _file_format.py so they stop tripping no-untyped-def. Import ChecksumWriter where needed. - Cast numpy returns whose dtype is exact but mypy widens to Any (argmin, astype, elementwise arithmetic): _kmeans.assign_l2, _kmeans.probe_scores_l2_monotone, _index._unpack_to_indices, _pq._preprocess_single, _ivfpq._preprocess_single. Plain typing.cast, no runtime overhead. - Lock three float-dtype assignments that np.where / division widen to float64: normalisation paths in _residual.add_batch, _pq.add_batch, _ivfpq.add_batch. Explicit .astype(np.float32) at the assignment. Infra: - Drop the deprecated numpy.typing.mypy_plugin from the mypy config. - CI: mypy step now fails on error instead of '|| true'. Result: 0 mypy errors, 190 tests still pass, ruff clean. --- .github/workflows/ci.yml | 4 ++-- pyproject.toml | 3 ++- snapvec/_fast.pyi | 28 ++++++++++++++++++++++++++++ snapvec/_file_format.py | 8 +++++++- snapvec/_index.py | 11 +++++++---- snapvec/_ivfpq.py | 14 +++++++------- snapvec/_kmeans.py | 8 ++++++-- snapvec/_pq.py | 14 +++++++------- snapvec/_residual.py | 8 ++++---- 9 files changed, 70 insertions(+), 28 deletions(-) create mode 100644 snapvec/_fast.pyi diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4be8a96..b5ebaeb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,8 +31,8 @@ jobs: - name: ruff check run: ruff check snapvec/ tests/ - - name: mypy (strict, warning-only for now) - run: mypy --strict snapvec/ || true + - name: mypy --strict + run: mypy --strict snapvec/ test: name: Test ${{ matrix.os }} / py${{ matrix.python-version }} diff --git a/pyproject.toml b/pyproject.toml index 0f000a5..4c0da76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,7 +70,8 @@ warn_return_any = true warn_unused_ignores = true disallow_untyped_defs = true disallow_any_generics = false -plugins = ["numpy.typing.mypy_plugin"] +# numpy.typing.mypy_plugin is deprecated as of NumPy 2.3; NDArray typing +# works without it on modern numpy. [[tool.mypy.overrides]] module = "tests.*" diff --git a/snapvec/_fast.pyi b/snapvec/_fast.pyi new file mode 100644 index 0000000..7aceae9 --- /dev/null +++ b/snapvec/_fast.pyi @@ -0,0 +1,28 @@ +"""Type stubs for the compiled Cython kernels (``_fast.pyx``). + +The real module is built from Cython and does not ship a ``.pyi`` +from the compiler; this stub lets ``mypy --strict`` see the same +Python-level shapes the Cython kernels expose to callers. +""" +from __future__ import annotations + +import numpy as np +from numpy.typing import NDArray + + +def adc_colmajor( + lut: NDArray[np.float32], + codes: NDArray[np.uint8], + scores: NDArray[np.float32], + parallel: bool = ..., +) -> None: ... + + +def fused_gather_adc( + all_codes: NDArray[np.uint8], + row_idx: NDArray[np.int64], + coarse_offsets: NDArray[np.float32], + lut: NDArray[np.float32], + scores: NDArray[np.float32], + parallel: bool = ..., +) -> None: ... diff --git a/snapvec/_file_format.py b/snapvec/_file_format.py index 1c6e825..81efc2f 100644 --- a/snapvec/_file_format.py +++ b/snapvec/_file_format.py @@ -30,6 +30,7 @@ import struct import zlib from pathlib import Path +from types import TracebackType from typing import IO, Callable @@ -81,7 +82,12 @@ def finalise(self) -> None: def __enter__(self) -> "ChecksumWriter": return self - def __exit__(self, exc_type, exc, tb) -> None: + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: if exc_type is None: self.finalise() diff --git a/snapvec/_index.py b/snapvec/_index.py index 1d35916..18428e6 100644 --- a/snapvec/_index.py +++ b/snapvec/_index.py @@ -45,13 +45,13 @@ import struct from pathlib import Path -from typing import Any +from typing import Any, cast import numpy as np from numpy.typing import NDArray from ._codebooks import get_codebook -from ._file_format import save_with_checksum_atomic, verify_checksum +from ._file_format import ChecksumWriter, save_with_checksum_atomic, verify_checksum from ._freezable import FreezableIndex from ._rotation import padded_dim, rht @@ -556,7 +556,7 @@ def save(self, path: str | Path) -> None: else: packed = _pack(self._indices, self._mse_bits) - def _write(f): + def _write(f: "ChecksumWriter") -> None: f.write(_MAGIC) f.write(struct.pack(" 1e-10, raw, 1.0) - units = arr / safe[:, None] + safe = np.where(raw > 1e-10, raw, 1.0).astype(np.float32) + units = (arr / safe[:, None]).astype(np.float32) norms = np.where(raw > 1e-10, raw, 0.0).astype(np.float32) if self.use_rht: padded = np.zeros((len(arr), self._pdim), dtype=np.float32) @@ -240,7 +240,7 @@ def _preprocess_single(self, q: NDArray[np.float32]) -> NDArray[np.float32]: padded[: self.dim] = q_unit rot = rht(padded[None, :], self.seed)[0] rot /= np.linalg.norm(rot) + 1e-12 - return rot.astype(np.float32) + return cast("NDArray[np.float32]", rot.astype(np.float32)) def _require_fitted(self) -> None: if not self._fitted: @@ -979,7 +979,7 @@ def save(self, path: str | Path) -> None: flags |= _FLAG_KEEP_FULL_PRECISION n = len(self._ids_by_row) - def _write(f): + def _write(f: "ChecksumWriter") -> None: f.write(_MAGIC) f.write( struct.pack( diff --git a/snapvec/_kmeans.py b/snapvec/_kmeans.py index 56d1c81..4eebd94 100644 --- a/snapvec/_kmeans.py +++ b/snapvec/_kmeans.py @@ -10,6 +10,7 @@ """ from __future__ import annotations +from typing import cast import numpy as np from numpy.typing import NDArray @@ -88,7 +89,7 @@ def assign_l2( ) -> NDArray[np.int64]: """Hard-assign every row in X to its nearest centroid (squared L2).""" d2 = (X ** 2).sum(1, keepdims=True) - 2 * X @ C.T + (C ** 2).sum(1)[None, :] - return d2.argmin(1) + return cast("NDArray[np.int64]", d2.argmin(1)) def probe_scores_l2_monotone( @@ -107,7 +108,10 @@ def probe_scores_l2_monotone( probe time gives slightly lower recall, especially with uneven cluster sizes. """ - return 2.0 * (coarse @ q) - (coarse ** 2).sum(1) + return cast( + "NDArray[np.float32]", + 2.0 * (coarse @ q) - (coarse ** 2).sum(1), + ) __all__ = [ diff --git a/snapvec/_pq.py b/snapvec/_pq.py index 6bbeef8..0913816 100644 --- a/snapvec/_pq.py +++ b/snapvec/_pq.py @@ -27,7 +27,7 @@ import struct from pathlib import Path -from typing import Any +from typing import Any, cast import numpy as np from numpy.typing import NDArray @@ -35,8 +35,8 @@ try: from ._fast import adc_colmajor except ImportError: - from ._fast_fallback import adc_colmajor # type: ignore[assignment] -from ._file_format import save_with_checksum_atomic, verify_checksum + from ._fast_fallback import adc_colmajor +from ._file_format import ChecksumWriter, save_with_checksum_atomic, verify_checksum from ._freezable import FreezableIndex from ._kmeans import kmeans_mse from ._rotation import padded_dim, rht @@ -150,8 +150,8 @@ def _preprocess( units = arr else: raw = np.linalg.norm(arr, axis=1) - safe = np.where(raw > 1e-10, raw, 1.0) - units = arr / safe[:, None] + safe = np.where(raw > 1e-10, raw, 1.0).astype(np.float32) + units = (arr / safe[:, None]).astype(np.float32) norms = np.where(raw > 1e-10, raw, 0.0).astype(np.float32) if self.use_rht: @@ -180,7 +180,7 @@ def _preprocess_single( padded[: self.dim] = q_unit rot = rht(padded[None, :], self.seed)[0] rot /= np.linalg.norm(rot) + 1e-12 - return rot.astype(np.float32) + return cast("NDArray[np.float32]", rot.astype(np.float32)) # ──────────────────────────────────────────────────────────────── # # training # @@ -372,7 +372,7 @@ def save(self, path: str | Path) -> None: flags |= _FLAG_USE_RHT n = len(self._ids) - def _write(f): + def _write(f: "ChecksumWriter") -> None: f.write(_MAGIC) f.write( struct.pack( diff --git a/snapvec/_residual.py b/snapvec/_residual.py index cc3cd90..126e640 100644 --- a/snapvec/_residual.py +++ b/snapvec/_residual.py @@ -31,7 +31,7 @@ from numpy.typing import NDArray from ._codebooks import get_codebook -from ._file_format import save_with_checksum_atomic, verify_checksum +from ._file_format import ChecksumWriter, save_with_checksum_atomic, verify_checksum from ._freezable import FreezableIndex from ._rotation import padded_dim, rht @@ -148,8 +148,8 @@ def add_batch( batch_norms = None # not stored in normalized mode else: raw_norms = np.linalg.norm(arr, axis=1) - safe = np.where(raw_norms > 1e-10, raw_norms, 1.0) - units = arr / safe[:, None] + safe = np.where(raw_norms > 1e-10, raw_norms, 1.0).astype(np.float32) + units = (arr / safe[:, None]).astype(np.float32) batch_norms = np.where(raw_norms > 1e-10, raw_norms, 0.0).astype(np.float32) pdim = self._pdim @@ -289,7 +289,7 @@ def save(self, path: str | Path) -> None: flags |= 1 n = len(self._ids) - def _write(f): + def _write(f: "ChecksumWriter") -> None: f.write(_MAGIC) f.write(struct.pack(" Date: Mon, 20 Apr 2026 16:26:24 +0200 Subject: [PATCH 2/3] fix: address PR #56 review -- eliminate unnecessary float32 copies Gemini flagged that the .astype(np.float32) calls I added to satisfy mypy introduce runtime copies on the hot path. Replace with typing.cast (pure type-level, zero cost) in five places: - _residual.add_batch, _pq.add_batch, _ivfpq.add_batch: wrap the np.where + division expressions with cast() instead of astype(). Use np.float32(1.0) / np.float32(0.0) as the False branch so np.where does not promote the result to float64 in the first place. - _pq._preprocess_single, _ivfpq._preprocess_single: drop the redundant .astype(np.float32) on 'rot'. rot is already float32 because it came out of a float32 padded buffer through rht(). cast() alone is enough. Net effect: removes a potential (N, dim) float32 copy per add_batch call and a (pdim,) copy per query on the preprocess path, with zero change to the semantics or the mypy strict guarantee. --- snapvec/_ivfpq.py | 14 ++++++++++---- snapvec/_pq.py | 14 ++++++++++---- snapvec/_residual.py | 9 ++++++--- 3 files changed, 26 insertions(+), 11 deletions(-) diff --git a/snapvec/_ivfpq.py b/snapvec/_ivfpq.py index 36df36f..69538c4 100644 --- a/snapvec/_ivfpq.py +++ b/snapvec/_ivfpq.py @@ -217,9 +217,15 @@ def _preprocess( norms = np.empty(0, dtype=np.float32) else: raw = np.linalg.norm(arr, axis=1) - safe = np.where(raw > 1e-10, raw, 1.0).astype(np.float32) - units = (arr / safe[:, None]).astype(np.float32) - norms = np.where(raw > 1e-10, raw, 0.0).astype(np.float32) + safe = cast( + "NDArray[np.float32]", + np.where(raw > 1e-10, raw, np.float32(1.0)), + ) + units = cast("NDArray[np.float32]", arr / safe[:, None]) + norms = cast( + "NDArray[np.float32]", + np.where(raw > 1e-10, raw, np.float32(0.0)), + ) if self.use_rht: padded = np.zeros((len(arr), self._pdim), dtype=np.float32) padded[:, : self.dim] = units @@ -240,7 +246,7 @@ def _preprocess_single(self, q: NDArray[np.float32]) -> NDArray[np.float32]: padded[: self.dim] = q_unit rot = rht(padded[None, :], self.seed)[0] rot /= np.linalg.norm(rot) + 1e-12 - return cast("NDArray[np.float32]", rot.astype(np.float32)) + return cast("NDArray[np.float32]", rot) def _require_fitted(self) -> None: if not self._fitted: diff --git a/snapvec/_pq.py b/snapvec/_pq.py index 0913816..8e898d9 100644 --- a/snapvec/_pq.py +++ b/snapvec/_pq.py @@ -150,9 +150,15 @@ def _preprocess( units = arr else: raw = np.linalg.norm(arr, axis=1) - safe = np.where(raw > 1e-10, raw, 1.0).astype(np.float32) - units = (arr / safe[:, None]).astype(np.float32) - norms = np.where(raw > 1e-10, raw, 0.0).astype(np.float32) + safe = cast( + "NDArray[np.float32]", + np.where(raw > 1e-10, raw, np.float32(1.0)), + ) + units = cast("NDArray[np.float32]", arr / safe[:, None]) + norms = cast( + "NDArray[np.float32]", + np.where(raw > 1e-10, raw, np.float32(0.0)), + ) if self.use_rht: padded = np.zeros((len(arr), self._pdim), dtype=np.float32) @@ -180,7 +186,7 @@ def _preprocess_single( padded[: self.dim] = q_unit rot = rht(padded[None, :], self.seed)[0] rot /= np.linalg.norm(rot) + 1e-12 - return cast("NDArray[np.float32]", rot.astype(np.float32)) + return cast("NDArray[np.float32]", rot) # ──────────────────────────────────────────────────────────────── # # training # diff --git a/snapvec/_residual.py b/snapvec/_residual.py index 126e640..dd648a7 100644 --- a/snapvec/_residual.py +++ b/snapvec/_residual.py @@ -25,7 +25,7 @@ import struct from pathlib import Path -from typing import Any +from typing import Any, cast import numpy as np from numpy.typing import NDArray @@ -148,8 +148,11 @@ def add_batch( batch_norms = None # not stored in normalized mode else: raw_norms = np.linalg.norm(arr, axis=1) - safe = np.where(raw_norms > 1e-10, raw_norms, 1.0).astype(np.float32) - units = (arr / safe[:, None]).astype(np.float32) + safe = cast( + "NDArray[np.float32]", + np.where(raw_norms > 1e-10, raw_norms, np.float32(1.0)), + ) + units = cast("NDArray[np.float32]", arr / safe[:, None]) batch_norms = np.where(raw_norms > 1e-10, raw_norms, 0.0).astype(np.float32) pdim = self._pdim From 443ff4951e508ab9cb389f035c7825a03beffe88 Mon Sep 17 00:00:00 2001 From: Jayson Steffens Date: Mon, 20 Apr 2026 16:35:15 +0200 Subject: [PATCH 3/3] fix: address PR #56 Copilot review Two follow-ups from the post-review scan on the earlier mypy cleanup: 1. Package the type stub. '_fast.pyi' and 'py.typed' were not picked up by the wheel build because setuptools.packages.find only grabs .py files by default, and MANIFEST.in listed neither. Add both to MANIFEST.in (sdist) and declare them in tool.setuptools.package-data (wheel). Verified the wheel now ships snapvec/_fast.pyi and snapvec/py.typed, so downstream mypy/pyright users get the types we added locally. 2. Defensive scalar promotion in probe_scores_l2_monotone. Modern numpy (NEP 50) keeps '2.0 * float32_array' as float32, but older numpy promoted it to float64, which would have made the cast() annotation a lie at runtime. Use np.float32(2.0) so the expression type-stays-put on both old and new numpy. No behaviour change on numpy 2.x. --- MANIFEST.in | 2 ++ pyproject.toml | 3 +++ snapvec/_kmeans.py | 6 +++++- 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/MANIFEST.in b/MANIFEST.in index 8e62988..3d539e0 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,5 +1,7 @@ include snapvec/_fast.pyx include snapvec/_fast.c +include snapvec/_fast.pyi +include snapvec/py.typed include CHANGELOG.md include LICENSE recursive-include tests *.py diff --git a/pyproject.toml b/pyproject.toml index 4c0da76..96b93cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,9 @@ docs = [ [tool.setuptools.packages.find] include = ["snapvec*"] +[tool.setuptools.package-data] +snapvec = ["py.typed", "*.pyi"] + [tool.pytest.ini_options] testpaths = ["tests"] pythonpath = ["."] diff --git a/snapvec/_kmeans.py b/snapvec/_kmeans.py index 4eebd94..1c8475d 100644 --- a/snapvec/_kmeans.py +++ b/snapvec/_kmeans.py @@ -108,9 +108,13 @@ def probe_scores_l2_monotone( probe time gives slightly lower recall, especially with uneven cluster sizes. """ + # np.float32(2.0) guards against NEP 50-era numpy upcasting a + # Python '2.0' scalar to float64 here; on numpy >= 2.0 this is a + # no-op, on older numpy it keeps the return dtype matching the + # annotation. return cast( "NDArray[np.float32]", - 2.0 * (coarse @ q) - (coarse ** 2).sum(1), + np.float32(2.0) * (coarse @ q) - (coarse ** 2).sum(1), )