more sophisticated type checking with numpy

[dcherian]

What is your issue?

This cool blogpost points out many improvements in numpy typing (NumType).

It also mentions stubtest which should be useful for our stubs, and also seems to be opt-in.

This would be a nice project for someone interested in type safety.

xref #10273

[adendek]

I’d like to start some initial investigation into this. If anyone else is interested in contributing or has specific ideas on the implementation, please let me know here! I’m happy to collaborate or step aside if someone else is already further along.

[dcherian]

Thank you! Go for it. Feel free to open a draft PR with questions if you need help

[kkollsga]

Tried running mypy.stubtest xarray locally on main to see what it would take. Got it working after a few fixes.

Getting stubtest to run

Out of the box, stubtest fails during the mypy build phase before it can even check stubs. The blockers were:

1. A handful of type errors in core code (~8) - mostly in daskmanager.py (override signatures) and generate_aggregations.py (wrong default value, missing method)
2. flox wasn't in the ignore_missing_imports list in pyproject.toml
3. Some now-unused type ignores that scipy-stubs made redundant

After fixing those, stubtest runs and produces actual stub validation output.

What stubtest found

Running on core modules (excluding tests):

python -m mypy.stubtest xarray.core xarray.backends xarray.namedarray \
    xarray.coding xarray.structure xarray.computation xarray.indexes \
    --mypy-config-file pyproject.toml

This produces ~6600 lines of stub discrepancies. Most are things like:

• _aggregations.py re-exports (Any, Callable, DataArray, Dataset) not matching runtime
• Missing or extra parameters in method signatures
• Type mismatches between stubs and runtime

Test files

Test files have additional issues - mainly pd.array() calls with pyarrow dtypes that pandas-stubs doesn't recognize, and some np.arange() calls that the stricter numpy typing doesn't like. These would need type ignores or could be excluded from stubtest.

Full stubtest output attached.

stubtest_output.txt

[dcherian]

Thanks @kkollsga ! We are happy to merge any PRs you can submit to fix these issues.

[kkollsga]

I've opened a PR to address this issue: #11124

Summary

The PR integrates stubtest into xarray's CI pipeline with:

1. CI Job - Runs stubtest on every PR (currently non-blocking with continue-on-error: true)
2. Allowlist - Comprehensive _stubtest/allowlist.txt covering:
   • TYPE_CHECKING imports (342 patterns)
   • Dynamically injected methods (Coarsen reductions)
   • Stdlib protocol differences (Mapping, Sequence)
   • scipy-stubs compatibility
3. Test infrastructure - ci_tests/ folder with pytest-based regression tests

Results

Metric	Before	After
Stubtest errors	1,084	0
Type ignores	+14	-8 net

Phased Rollout

The stubtest job is set to continue-on-error: true (Phase 1: informational). Once you are confident with the changes, it can be made required by changing to continue-on-error: false.

[kkollsga]

Ran a deeper analysis on the ~6000+ stubtest errors to understand what they actually represent.

Results

The errors fall into predictable categories - none of which are actual type annotation bugs:

• TYPE_CHECKING imports (~60%): Types like Self, Dims, DataArray imported under if TYPE_CHECKING: blocks to avoid circular imports. These exist for static analysis but not at runtime - intentional and correct.

• Protocol/ABC differences (~25%): Standard library ABCs (Mapping, MutableMapping, Sequence) have slight signature differences between runtime and stubs. Not xarray issues.

• Accessor patterns (~10%): Plot accessors use descriptors that stubtest doesn't fully understand. Working as intended.

• Other expected patterns (~5%): CFTimeIndex read-only properties, numpy protocol methods, etc.

What this means

xarray's type annotations are in good shape - no source code fixes were needed. The updated PR #11124 is infrastructure-only: a CI job and an allowlist covering these known patterns.

The allowlist can be updated as the project evolves - patterns can be removed as upstream stubs improve, or added if new acceptable differences arise. Future real type regressions will now be caught in CI.