TYP: Improved ndarray augmented assignment operators
#29862
+106
−215
Conversation
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
|
@MarcoGorelli do you think we could also apply this trick to get rid of the other arithmetic dunder overrides in |
jorenham
commented
Oct 2, 2025
Comment on lines
3437
to
+3442
| @overload | ||
| def __iadd__(self: NDArray[integer], other: _ArrayLikeInt_co, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ... | ||
| def __iadd__(self: _ComplexFloatingArrayT, other: _ArrayLikeComplex_co, /) -> _ComplexFloatingArrayT: ... | ||
| @overload | ||
| def __iadd__(self: NDArray[floating], other: _ArrayLikeFloat_co, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ... | ||
| def __iadd__(self: _InexactArrayT, other: _ArrayLikeFloat_co, /) -> _InexactArrayT: ... | ||
| @overload | ||
| def __iadd__(self: NDArray[complexfloating], other: _ArrayLikeComplex_co, /) -> ndarray[_ShapeT_co, _DTypeT_co]: ... | ||
| def __iadd__(self: _NumberArrayT, other: _ArrayLikeInt_co, /) -> _NumberArrayT: ... |
There was a problem hiding this comment.
I had to flip the overload order here to appease mypy. It started complaining about non-existent overlap once I put the _*ArrayT typars in, demonstrating that the error is a false-positive.
|
Diff from mypy_primer, showing the effect of this PR on type check results on a corpus of open source code: colour (https://github.com/colour-science/colour)
- colour/colorimetry/tristimulus_values.py:458: error: Argument 1 to "__imul__" of "ndarray" has incompatible type "object"; expected "_SupportsArray[dtype[numpy.bool[builtins.bool] | integer[Any] | floating[Any]]] | _NestedSequence[_SupportsArray[dtype[numpy.bool[builtins.bool] | integer[Any] | floating[Any]]]] | float | _NestedSequence[float]" [arg-type]
+ colour/colorimetry/tristimulus_values.py:458: error: No overload variant of "__imul__" of "ndarray" matches argument type "object" [call-overload]
+ colour/colorimetry/tristimulus_values.py:458: note: Possible overload variants:
+ colour/colorimetry/tristimulus_values.py:458: note: def __imul__(self, _SupportsArray[dtype[numpy.bool[builtins.bool] | integer[Any] | floating[Any]]] | _NestedSequence[_SupportsArray[dtype[numpy.bool[builtins.bool] | integer[Any] | floating[Any]]]] | float | _NestedSequence[float], /) -> ndarray[tuple[Any, ...], dtype[floating[Any]]]
+ colour/colorimetry/tristimulus_values.py:458: note: def __imul__(self, _SupportsArray[dtype[numpy.bool[builtins.bool] | integer[Any]]] | _NestedSequence[_SupportsArray[dtype[numpy.bool[builtins.bool] | integer[Any]]]] | int | _NestedSequence[int], /) -> ndarray[tuple[Any, ...], dtype[floating[Any]]]
|
|
@MarcoGorelli Ping. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
By using type-restricted method-bound type parameters for
selfin thendarray.__i{op}__"augmented assignment" operator methods, they'll now covariantly specialize, and side-step the need for higher-kinded types (sorry about the jargon). I.e., subtypes ofndarrayno longer need to override the__i{op}__methods.This allows us to get rid of a bunch of duplication in e.g.
np.ma.MaskedArray, and users with their ownndaraysubclasses are now also able to do the same.This also addresses the issue described in python/mypy#19171 more rigourously than #29092, since the
self-bound type-parameters are now slightly wider. For example thendarray.__add__overload that accepts "integer-like array-likes" (_ArrayLikeInt_co) now applies to allnp.numberarrays, instead of just thenp.integerarrays. For floating-like array-likes thenp.floatingrestriction is now similarly widened tonp.inexact. I believe this resolves the remaining class ofissues described in python/mypy#19171 (comment) by @david-zwicker.