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/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 0f000a5..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 = ["."] @@ -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.*" 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] - 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( diff --git a/snapvec/_kmeans.py b/snapvec/_kmeans.py index 56d1c81..1c8475d 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,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__ = [ diff --git a/snapvec/_pq.py b/snapvec/_pq.py index 6bbeef8..8e898d9 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,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( diff --git a/snapvec/_residual.py b/snapvec/_residual.py index cc3cd90..dd648a7 100644 --- a/snapvec/_residual.py +++ b/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("