chore: fix all mypy strict errors, promote to hard CI gate

.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 }}

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

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 = ["."]
@@ -70,7 +73,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.*"

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: ...

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

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("<IIIIII", _VERSION, self.dim, self.bits, self.seed, n, flags))
             f.write(struct.pack("<I", len(packed)))
@@ -800,7 +800,10 @@ def _unpack(
         return result
     if bits == 3 and not legacy_3bit:
         arr = np.frombuffer(data, dtype=np.uint8)
-        return _unpack_3bit_tight(arr)[:total].reshape(n_rows, n_cols).copy()
+        return cast(
+            "NDArray[np.uint8]",
+            _unpack_3bit_tight(arr)[:total].reshape(n_rows, n_cols).copy(),
+        )
     # Byte-aligned path: 2-bit, 4-bit, and legacy 3-bit (v1/v2)
     ipb = 8 // bits
     mask = (1 << bits) - 1

snapvec/_ivfpq.py

@@ -34,16 +34,16 @@
 import warnings
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 from numpy.typing import NDArray
 
 try:
     from ._fast import fused_gather_adc
 except ImportError:
-    from ._fast_fallback import fused_gather_adc  # type: ignore[assignment]
-from ._file_format import save_with_checksum_atomic, verify_checksum
+    from ._fast_fallback import fused_gather_adc
+from ._file_format import ChecksumWriter, save_with_checksum_atomic, verify_checksum
 from ._freezable import FreezableIndex
 from ._kmeans import assign_l2, kmeans_mse, probe_scores_l2_monotone
 from ._rotation import padded_dim, rht
@@ -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)
-            units = arr / safe[:, None]
-            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 rot.astype(np.float32)
+        return cast("NDArray[np.float32]", rot)
 
     def _require_fitted(self) -> None:
         if not self._fitted:
@@ -979,7 +985,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(

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,14 @@ 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)
+    # 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]",
+        np.float32(2.0) * (coarse @ q) - (coarse ** 2).sum(1),
+    )
 
 
 __all__ = [

snapvec/_pq.py

@@ -27,16 +27,16 @@
 
 import struct
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 from numpy.typing import NDArray
 
 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,9 +150,15 @@ 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]
-            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 rot.astype(np.float32)
+        return cast("NDArray[np.float32]", rot)
 
     # ──────────────────────────────────────────────────────────────── #
     # training                                                          #
@@ -372,7 +378,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(

snapvec/_residual.py

@@ -25,13 +25,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
 
@@ -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)
-            units = arr / safe[:, None]
+            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
@@ -289,7 +292,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("<IIIIIIII", _VERSION, self.dim, self.b1,
                                 self.b2, self.seed, n, flags, self._pdim))