TYP: sum and prod shape-typing and improved dtypes

[jorenham]

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]

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.

[github-actions]

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]

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]

Thanks Joren.