config: add quax to the mypy pre-commit hook - #806
Merged
Conversation
The mypy hook's additional_dependencies deliberately excludes quax, quaxed, unxt, coordinax, and jax to keep the isolated hook environment fast -- everything they touch resolves to Any there. That's why galax.utils._jax.quaxify wrapped quax.quaxify in an explicit cast: without quax installed, decorating with it directly makes the decorated function untyped, which is what #796's review comment ran into. Installing quax specifically (not the rest -- that pulls in the whole JAX stack and defeats the point of the minimal hook) lets mypy follow real types through every quax-touched call chain for the first time. That surfaced 191 pre-existing errors across 60 files that were previously masked, mostly: - `no-any-return` from jax's own functions (jax.jit, jax.grad, jax.lax.scan, .min()/.max()) being untyped under this project's ignore_missing_imports policy for jax.* -- fixed with `# type: ignore[no-any-return]`, matching the one existing precedent in this codebase (coordinates/_src/base.py). - redundant casts that predated quax being resolvable, now removed, including the galax.utils._jax.quaxify wrapper itself. - `override`/`misc` from equinox's AbstractVar/AbstractClassVar + property pattern and plum's multiple-dispatch redefinitions -- ignored with the matching error code. - two real fixes: _error_if_not_all_constant_parameters in the galpy interop is now generic so it preserves the caller's specific potential subtype instead of widening to AbstractPotential; and NBodyField._call used `self.eps.ustrip(...)`/`self.masses.ustrip(...)` as method calls, which would crash if a bare Array (not a Quantity) were passed for either field -- switched to the polymorphic `u.ustrip(AllowValue, ...)` free function already used two lines above for the same purpose. - `integrate_field`'s `@ft.partial(eqx.filter_jit)` decorator carried zero bound arguments, making it a no-op wrapper that also breaks mypy's ability to see through functools.partial's ParamSpec; changed to bare `@eqx.filter_jit`, matching the pattern already used elsewhere in this codebase. Verified with the actual pre-commit hook (not just a local mypy invocation): `pre-commit run --all-files` is clean, the full test suite passes, and doctests on every touched file pass. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #806 +/- ##
==========================================
- Coverage 95.75% 95.73% -0.03%
==========================================
Files 159 159
Lines 6055 6089 +34
==========================================
+ Hits 5798 5829 +31
- Misses 257 260 +3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Followup to a review comment on #805/#796:
quaxifyis wrapped ingalax.utils._jax/galax.coordinates._src.shapewith an explicitcast, becausequaxis deliberately excluded from the mypy pre-commit hook's isolated environment (onlydataclassish,optype, andplum-dispatchare installed there, to keep it fast). Withoutquaxresolvable, decorating withquax.quaxifydirectly makes the decorated function untyped — so the cast isn't a workaround forquax's own typing, it's a workaround for this hook's minimal environment.This PR adds
quax(onlyquax— notquaxed/unxt/coordinax/jax, which would pull in the whole JAX stack and defeat the point of the minimal hook) toadditional_dependencies, and removes the now-genuinely-redundantquaxifywrapper.Doing so lets mypy follow real types through every
quax-touched call chain for the first time, which surfaced 191 pre-existing errors across 60 files that were previously masked. Breakdown:no-any-return(bulk of it) — fromjax's own functions (jax.jit,jax.grad,jax.lax.scan,.min()/.max()) being untyped under this project'signore_missing_importspolicy forjax.*. Fixed with# type: ignore[no-any-return], matching the one existing precedent in this codebase (coordinates/_src/base.py).quaxbeing resolvable — removed, including thequaxifywrapper itself.override/miscfrom equinox'sAbstractVar/AbstractClassVar+ property-override pattern, and plum's multiple-dispatch redefinitions — ignored with the matching error code, consistent with how this codebase already handles the same pattern elsewhere._error_if_not_all_constant_parametersin the galpy interop is now generic ([PT: gp.AbstractPotential](pot: PT) -> PT) so it preserves the caller's specific potential subtype instead of widening every caller toAbstractPotential.NBodyField's force calculation calledself.eps.ustrip(...)/self.masses.ustrip(...)as method calls, but both fields are typedQuantity | Array— a bareArrayhas no.ustripmethod, so this would crash at runtime for that case. Switched to the polymorphicu.ustrip(AllowValue, ...)free function already used two lines above for the same purpose.integrate_field's@ft.partial(eqx.filter_jit)decorator had zero bound arguments, making it a no-op wrapper that also breaks mypy's ability to seeeqx.filter_jit'sParamSpec-preserving signature throughfunctools.partial. Simplified to bare@eqx.filter_jit, matching the pattern already used elsewhere in this codebase.Verification
pre-commit run --all-files(the actual hook config, not a local approximation) is clean.test_mockstreamgenerator.py::test_second_deriv, was confirmed to already fail identically on unmodifiedmain; unrelated to this change, deselected for the full-suite verification run.)🤖 Generated with Claude Code