TYP: Type MaskedArray.__{iadd,isub,imul,itruediv,ifloordiv,pow}__

MarcoGorelli · MarcoGorelli · commit 71b1258e845c · 2025-05-16T14:53:34.000+01:00
diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi
@@ -17,22 +17,34 @@ from numpy import (
     amax,
     amin,
     bool_,
+    complex128,
+    complexfloating,
+    datetime64,
     dtype,
     expand_dims,
     float64,
+    floating,
     generic,
     int_,
     intp,
     ndarray,
+    object_,
+    signedinteger,
+    timedelta64,
+    unsignedinteger,
 )
 from numpy._globals import _NoValueType
 from numpy._typing import (
     ArrayLike,
     NDArray,
     _ArrayLike,
     _ArrayLikeBool_co,
+    _ArrayLikeComplex_co,
+    _ArrayLikeFloat_co,
     _ArrayLikeInt,
     _ArrayLikeInt_co,
+    _ArrayLikeTD64_co,
+    _ArrayLikeUInt_co,
     _DTypeLikeBool,
     _IntLike_co,
     _ScalarLike_co,
@@ -456,12 +468,170 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     def __rfloordiv__(self, other): ...
     def __pow__(self, other, mod: None = None, /): ...
     def __rpow__(self, other, mod: None = None, /): ...
-    def __iadd__(self, other): ...
-    def __isub__(self, other): ...
-    def __imul__(self, other): ...
-    def __ifloordiv__(self, other): ...
-    def __itruediv__(self, other): ...
-    def __ipow__(self, other): ...
+
+    # Keep in sync with `ndarray.__iadd__`, except that `_MaskedArray[unsignedinteger]` does not accept
+    # _IntLake_co for `other`.
+    @overload
+    def __iadd__(
+        self: _MaskedArray[np.bool], other: _ArrayLikeBool_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __iadd__(
+        self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __iadd__(
+        self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __iadd__(
+        self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __iadd__(
+        self: _MaskedArray[complexfloating], other: _ArrayLikeComplex_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __iadd__(
+        self: _MaskedArray[timedelta64 | datetime64], other: _ArrayLikeTD64_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __iadd__(
+        self: _MaskedArray[object_], other: Any, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+
+    # Keep in sync with `ndarray.__isub__`, except that `_MaskedArray[unsignedinteger]` does not accept
+    # _IntLike_co for `other`.
+    @overload
+    def __isub__(
+        self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __isub__(
+        self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __isub__(
+        self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __isub__(
+        self: _MaskedArray[complexfloating], other: _ArrayLikeComplex_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __isub__(
+        self: _MaskedArray[timedelta64 | datetime64], other: _ArrayLikeTD64_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __isub__(
+        self: _MaskedArray[object_], other: Any, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+
+    # Keep in sync with `ndarray.__imul__`, except that `_MaskedArray[unsignedinteger]` does not accept
+    # _IntLike_co for `other`.
+    @overload
+    def __imul__(
+        self: _MaskedArray[np.bool], other: _ArrayLikeBool_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __imul__(
+        self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __imul__(
+        self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __imul__(
+        self: _MaskedArray[float64], other: _ArrayLikeFloat_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __imul__(
+        self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __imul__(
+        self: _MaskedArray[complex128], other: _ArrayLikeComplex_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __imul__(
+        self: _MaskedArray[complexfloating], other: _ArrayLikeComplex_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __imul__(
+        self: _MaskedArray[timedelta64], other: _ArrayLikeFloat_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __imul__(
+        self: _MaskedArray[object_], other: Any, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+
+    # Keep in sync with `ndarray.__ifloordiv__`, except that `_MaskedArray[unsignedinteger]` does not accept
+    # _IntLike_co for `other`.
+    @overload
+    def __ifloordiv__(
+        self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __ifloordiv__(
+        self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __ifloordiv__(
+        self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __ifloordiv__(
+        self: _MaskedArray[timedelta64], other: _ArrayLikeInt, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __ifloordiv__(
+        self: _MaskedArray[object_], other: Any, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+
+    # Keep in sync with `ndarray.__itruediv__`, except that `_MaskedArray[unsignedinteger]` does not accept
+    # _IntLike_co for `other`.
+    @overload
+    def __itruediv__(
+        self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __itruediv__(
+        self: _MaskedArray[complexfloating],
+        other: _ArrayLikeComplex_co,
+        /,
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __itruediv__(
+        self: _MaskedArray[timedelta64], other: _ArrayLikeInt, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __itruediv__(
+        self: _MaskedArray[object_], other: Any, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+
+    # Keep in sync with `ndarray.__ipow__`, except that `_MaskedArray[unsignedinteger]` does not accept
+    # _IntLike_co for `other`.
+    @overload
+    def __ipow__(
+        self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __ipow__(
+        self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __ipow__(
+        self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __ipow__(
+        self: _MaskedArray[complexfloating], other: _ArrayLikeComplex_co, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
+    @overload
+    def __ipow__(
+        self: _MaskedArray[object_], other: Any, /
+    ) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
     @property  # type: ignore[misc]
     def imag(self: _HasDTypeWithRealAndImag[object, _ScalarT], /) -> MaskedArray[_ShapeT_co, dtype[_ScalarT]]: ...
     get_imag: Any
diff --git a/numpy/typing/tests/data/fail/ma.pyi b/numpy/typing/tests/data/fail/ma.pyi
@@ -1,7 +1,17 @@
+from typing import TypeAlias, TypeVar
+
 import numpy as np
 import numpy.typing as npt
+from numpy import dtype, generic
+from numpy._typing import _Shape
+
+_ScalarT = TypeVar("_ScalarT", bound=generic)
+MaskedArray: TypeAlias = np.ma.MaskedArray[_Shape, dtype[_ScalarT]]
 
 MAR_1d_f8: np.ma.MaskedArray[tuple[int], np.dtype[np.float64]]
+MAR_b: MaskedArray[np.bool]
+MAR_c: MaskedArray[np.complex128]
+MAR_td64: MaskedArray[np.timedelta64]
 
 AR_b: npt.NDArray[np.bool]
 
@@ -127,4 +137,8 @@ np.ma.allclose(MAR_1d_f8, [1,2,3], atol='.5')  # type: ignore[arg-type]
 
 MAR_1d_f8.__setmask__('mask')  # type: ignore[arg-type]
 
+MAR_b *= 2  # type: ignore[arg-type]
+MAR_c //= 2  # type: ignore[misc]
+MAR_td64 **= 2  # type: ignore[misc]
+
 MAR_1d_f8.swapaxes(axis1=1, axis2=0)  # type: ignore[call-arg]
diff --git a/numpy/typing/tests/data/pass/ma.py b/numpy/typing/tests/data/pass/ma.py
