chore: fix all mypy strict errors, promote to hard CI gate - #56
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request improves type safety and compatibility by adding type stubs for Cython kernels, updating configuration for NumPy 2.3, and refining type hints across several modules. Review feedback highlights opportunities to optimize performance in numerical computation paths by using typing.cast or astype(..., copy=False) to avoid redundant array copies and preventing implicit promotion to float64 by using np.float32 literals.
| safe = np.where(raw > 1e-10, raw, 1.0).astype(np.float32) | ||
| units = (arr / safe[:, None]).astype(np.float32) |
There was a problem hiding this comment.
Using .astype(np.float32) on the result of a division or np.where creates an unnecessary runtime copy of the array. Additionally, using the Python float literal 1.0 in np.where promotes the intermediate result to float64. Since this code is in a performance-sensitive path, consider using np.float32(1.0) to avoid promotion and typing.cast to satisfy the type checker without the runtime cost.
| safe = np.where(raw > 1e-10, raw, 1.0).astype(np.float32) | |
| units = (arr / safe[:, None]).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]) |
| 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)) |
There was a problem hiding this comment.
Redundant use of both cast and .astype(np.float32). Since rot is already a float32 array, astype creates an unnecessary copy. You can use cast alone to satisfy the type checker.
| return cast("NDArray[np.float32]", rot.astype(np.float32)) | |
| return cast("NDArray[np.float32]", rot) |
| safe = np.where(raw > 1e-10, raw, 1.0).astype(np.float32) | ||
| units = (arr / safe[:, None]).astype(np.float32) |
There was a problem hiding this comment.
Using .astype(np.float32) here introduces unnecessary runtime overhead. Using np.float32(1.0) in np.where avoids promotion to float64, and typing.cast can be used to satisfy the type checker without creating a copy of the array.
| safe = np.where(raw > 1e-10, raw, 1.0).astype(np.float32) | |
| units = (arr / safe[:, None]).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]) |
| 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)) |
There was a problem hiding this comment.
Redundant use of both cast and .astype(np.float32). Since rot is already a float32 array, astype creates an unnecessary copy. You can use cast alone to satisfy the type checker.
| return cast("NDArray[np.float32]", rot.astype(np.float32)) | |
| return cast("NDArray[np.float32]", rot) |
| safe = np.where(raw_norms > 1e-10, raw_norms, 1.0).astype(np.float32) | ||
| units = (arr / safe[:, None]).astype(np.float32) |
There was a problem hiding this comment.
Using .astype(np.float32) here creates unnecessary runtime copies. To avoid promotion to float64, use np.float32(1.0) in np.where. To satisfy the type checker efficiently without an extra copy, you can use .astype(np.float32, copy=False).
| safe = np.where(raw_norms > 1e-10, raw_norms, 1.0).astype(np.float32) | |
| units = (arr / safe[:, None]).astype(np.float32) | |
| safe = np.where(raw_norms > 1e-10, raw_norms, np.float32(1.0)).astype(np.float32, copy=False) | |
| units = (arr / safe[:, None]).astype(np.float32, copy=False) |
There was a problem hiding this comment.
Pull request overview
This PR resolves remaining mypy --strict issues across snapvec/ and upgrades mypy from a warning-only check to a required CI gate, improving static type safety for the core index implementations and build/runtime pathways.
Changes:
- Add a stub for the compiled
_fastextension and remove now-unneeded typing ignores in fallback imports. - Tighten/repair typing around save-path closures (
ChecksumWriter), NumPy return types (cast()), and float32 normalization paths. - Update infra: remove deprecated NumPy mypy plugin usage and make
mypy --stricta hard CI requirement.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
snapvec/_fast.pyi |
Adds type stubs for compiled Cython kernels so strict mypy can type-check call sites. |
snapvec/_file_format.py |
Fully types ChecksumWriter.__exit__ to satisfy strict mypy. |
snapvec/_index.py |
Adds closure param typing for save writer and casts a NumPy-return path to keep types precise. |
snapvec/_pq.py |
Removes fallback import ignore, adds save-writer typing, and uses cast() for NumPy-typed returns. |
snapvec/_ivfpq.py |
Mirrors PQ typing adjustments for fallback import, preprocessing, and save-writer closure typing. |
snapvec/_residual.py |
Ensures normalization stays float32 and adds typed save-writer closure signature. |
snapvec/_kmeans.py |
Adds cast() where NumPy returns Any under strict typing. |
pyproject.toml |
Removes deprecated NumPy mypy plugin config (commented rationale). |
.github/workflows/ci.yml |
Promotes mypy --strict to a required CI check (no longer ` |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| """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. | ||
| """ |
There was a problem hiding this comment.
This stub is added to satisfy type-checking of the compiled extension, but the repo packaging config doesn’t currently declare .pyi / py.typed as package data (e.g., MANIFEST.in doesn’t include them). That means the stub may be missing from sdists/wheels, defeating the purpose for downstream type checkers. Ensure _fast.pyi (and py.typed, if intended) are included in the built distribution via MANIFEST.in and/or tool.setuptools.package-data / include-package-data configuration.
| return 2.0 * (coarse @ q) - (coarse ** 2).sum(1) | ||
| return cast( | ||
| "NDArray[np.float32]", | ||
| 2.0 * (coarse @ q) - (coarse ** 2).sum(1), |
There was a problem hiding this comment.
probe_scores_l2_monotone is annotated/cast to return NDArray[np.float32], but the expression 2.0 * (coarse @ q) - ... will upcast to float64 at runtime because 2.0 is a Python float. This makes the cast incorrect and can increase memory/compute. Use a float32 scalar (e.g., np.float32(2.0)) and/or explicitly .astype(np.float32) on the result, or change the return type annotation to float64 to match reality.
| 2.0 * (coarse @ q) - (coarse ** 2).sum(1), | |
| (np.float32(2.0) * (coarse @ q) - (coarse ** 2).sum(1)).astype(np.float32), |
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.
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.
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.
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.
5eed3db to
443ff49
Compare
Summary
Closes out the long-standing "17 mypy errors, warning-only in CI" item
from the original audit. After this PR, `mypy --strict` is a required
CI check, not a silent pass.
Changes
and `fused_gather_adc` signatures so mypy sees the compiled module.
Removes the `import-not-found` errors and the now-unused
`type: ignore` comments on the `_fast_fallback` branch in `_pq.py`
/ `_ivfpq.py`.
(`_index.py`, `_pq.py`, `_residual.py`, `_ivfpq.py`) and
`ChecksumWriter.exit` now have full signatures. Each file
that uses `_write` imports `ChecksumWriter` directly (not
`TYPE_CHECKING`-guarded, since it is used at save time).
arithmetic on typed arrays. mypy widens some of those to `Any`
under the (now-removed) numpy plugin; explicit `cast` keeps the
source-of-truth annotations honest without runtime cost.
paths (`_residual`, `_pq`, `_ivfpq`) explicitly cast via
`.astype(np.float32)` so the variable type matches the annotation.
of NumPy 2.3).
Test plan