Skip to content

TYP: sum and prod shape-typing and improved dtypes#31160

Merged
charris merged 6 commits into
numpy:mainfrom
jorenham:typing/sum
Apr 6, 2026
Merged

TYP: sum and prod shape-typing and improved dtypes#31160
charris merged 6 commits into
numpy:mainfrom
jorenham:typing/sum

Conversation

@jorenham

@jorenham jorenham commented Apr 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

This follows #31154, and uses a similar overload pattern as for np[.ndarray].{mean,std,var}, but sum and prod differ in that they don't upcast integer to float64, and instead upcast bool_ to int_, so they require additional overloads.

The sum and prod contracts are almost identical except for timedelta64; which is accepted sum but not by prod.

This also adds specialized overrides for the sum and prod methods of the number, bool, and timedelta64 (only for sum) scalar types, which might help in case you e.g. .sum() on float64 | NDArray[float64] input, i.e. when you're not sure if you're summing a scalar or an array.

AI Disclosure

N/A

@jorenham jorenham added this to the 2.5.0 Release milestone Apr 5, 2026
@github-actions

This comment has been minimized.

Sign in to view
@jorenham

jorenham commented Apr 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

RE: mypy_primer

  • jax: changed mypy error message
  • staticframe:
    • static_frame/core/series.py:3026: changed mypy error code
    • static_frame/core/bus.py:979: a comparison between a np.number and int | None, which will raise an error at runtime if the RHS is None, so mypy is right to complain
  • pandas:
    • pandas/core/groupby/generic.py:1281: This is actually a typing issue in numpy.tile, which currently doesn't accept np.int_ as second argument, even though it should. I'll fix this ASAP.
    • all other pandas errors are correct, because it's indeed type-unsafe to assign a np.int_ to a python int, which is what mypy rightly complains about
  • colour:
    • colour/phenomena/sky/wilkie2021.py:927-930: These are errors in colour itself, because it's trying to assign ndarray instances to dataclass fields of type float, which is type-unsafe.
    • colour/recovery/meng2015.py:191: this is caused by a known mypy bug, where it selects the first overload in case the dtype is Any, which is incorrect according to the typing spec. I think I can work around this though.
  • xarray: This error is a mypy issue, which can be solved by using the --allow-redefinition-new mypy flag.
  • optuna: it returns the product of a 2d array with axis=1 from a function annotated to return float, so the return annotation is incorrect
  • scipy: the function returns a tuple with np.int_ as fourth value, but the return type for that value is a Python int, which is incompatible, so the mypy error is accurate.
  • numpy-stl: again that mypy overload bug that I'll try to find a workaround for.

So to summarize, there are two action points for me:

  • There'a a bug in the numpy.tile stubs. I'll address this in a separate PR. Done: see TYP: tile: accept numpy scalars and arrays as second argument #31161
  • A mypy bug related to @overload is causing issues if the dtype is Any. I'll include a workaround for this in this PR. Done: I managed to work around this by changing the order of overloads.

@jorenham jorenham mentioned this pull request Apr 5, 2026
Merged
@github-actions

github-actions Bot commented Apr 5, 2026

Copy link
Copy Markdown

Diff from mypy_primer, showing the effect of this PR on type check results on a corpus of open source code:

jax (https://github.com/google/jax)
- jax/_src/pallas/pipelining/schedulers.py:412: error: Argument 3 to "check_async_done" has incompatible type "signedinteger[_64Bit]"; expected "int | Array"  [arg-type]
+ jax/_src/pallas/pipelining/schedulers.py:412: error: Argument 3 to "check_async_done" has incompatible type "signedinteger[_32Bit | _64Bit]"; expected "int | Array"  [arg-type]

colour (https://github.com/colour-science/colour)
+ colour/phenomena/sky/wilkie2021.py:927: error: Argument "theta" to "SkyParameters_Wilkie2021" has incompatible type "ndarray[tuple[Any, ...], dtype[Any]]"; expected "float"  [arg-type]
+ colour/phenomena/sky/wilkie2021.py:928: error: Argument "gamma" to "SkyParameters_Wilkie2021" has incompatible type "ndarray[tuple[Any, ...], dtype[Any]]"; expected "float"  [arg-type]
+ colour/phenomena/sky/wilkie2021.py:929: error: Argument "shadow" to "SkyParameters_Wilkie2021" has incompatible type "ndarray[tuple[Any, ...], dtype[Any]]"; expected "float"  [arg-type]
+ colour/phenomena/sky/wilkie2021.py:930: error: Argument "zero" to "SkyParameters_Wilkie2021" has incompatible type "ndarray[tuple[Any, ...], dtype[Any]]"; expected "float"  [arg-type]

pandas (https://github.com/pandas-dev/pandas)
+ pandas/io/parsers/base_parser.py:516: error: Incompatible types in assignment (expression has type "signedinteger[_32Bit | _64Bit]", variable has type "int")  [assignment]
+ pandas/io/parsers/base_parser.py:518: error: Incompatible types in assignment (expression has type "signedinteger[_32Bit | _64Bit]", variable has type "int")  [assignment]
+ pandas/core/internals/blocks.py:1279: error: Argument 2 to "setitem_datetimelike_compat" has incompatible type "signedinteger[_32Bit | _64Bit]"; expected "int"  [arg-type]
+ pandas/core/indexes/base.py:5773: error: Argument 2 to "setitem_datetimelike_compat" has incompatible type "signedinteger[_32Bit | _64Bit]"; expected "int"  [arg-type]
+ pandas/core/resample.py:2623: error: Incompatible types in assignment (expression has type "signedinteger[_32Bit | _64Bit]", variable has type "int")  [assignment]

optuna (https://github.com/optuna/optuna)
+ optuna/importance/_fanova/_tree.py:311: error: Incompatible return value type (got "ndarray[Any, Any]", expected "float")  [return-value]

numpy-stl (https://github.com/WoLpH/numpy-stl)
+ stl/base.py:448: error: Incompatible types in assignment (expression has type "ndarray[tuple[int, int], dtype[float64]]", variable has type "ndarray[tuple[int, int], dtype[floating[_32Bit]]]")  [assignment]

@jorenham

jorenham commented Apr 5, 2026

Copy link
Copy Markdown
Member Author

The remaining primer errors are, except for the numpy-stl one, all correct, and a consequence of incorrect annotations in these packages.

The numpy-stl error is caused by that mypy overload bug, where it selects the wrong ndarray.__mul__ overload when the dtype is Any. So that's out of scope for this PR.

@charris charris merged commit 33ea714 into numpy:main Apr 6, 2026
14 checks passed
@charris

charris commented Apr 6, 2026

Copy link
Copy Markdown
Member

Thanks Joren.

@jorenham jorenham deleted the typing/sum branch April 6, 2026 09:16
This was referenced Apr 6, 2026
Merged
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

41 - Static typing

Projects

None yet

Milestone

2.5.0 Release

Development

Successfully merging this pull request may close these issues.

2 participants

@jorenham @charris