TYP: Type MaskedArray.{trace,round,cumsum,cumprod}

numpy · jorenham · Jul 2, 2025 · Jun 26, 2025 · Jul 2, 2025 · Jul 2, 2025
commit 830d1ec7c888de55f7b9d0ec537f5ef4ced81e64
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -1741,6 +1741,7 @@ class _ArrayOrScalarCommon:
     @overload
     def argmin(self, /, axis: SupportsIndex | None = None, *, out: _BoolOrIntArrayT, keepdims: builtins.bool = False) -> _BoolOrIntArrayT: ...
 
+    # Keep in sync with `MaskedArray.round`
     @overload  # out=None (default)
     def round(self, /, decimals: SupportsIndex = 0, out: None = None) -> Self: ...
     @overload  # out=ndarray
@@ -2385,6 +2386,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
         stable: bool | None = ...,
     ) -> None: ...
 
+    # Keep in sync with `MaskedArray.trace`
     @overload
     def trace(
         self,  # >= 2D array

diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi
@@ -46,6 +46,7 @@ from numpy import (
 from numpy._globals import _NoValueType
 from numpy._typing import (
     ArrayLike,
+    DTypeLike,
     NDArray,
     _32Bit,
     _64Bit,
@@ -1168,18 +1169,70 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     ) -> _ArrayT: ...
 
     def nonzero(self) -> tuple[_Array1D[intp], ...]: ...
-    def trace(self, offset=..., axis1=..., axis2=..., dtype=..., out=...): ...
+
+    # Keep in sync with `ndarray.trace`
+    @overload
+    def trace(
+        self,  # >= 2D MaskedArray
+        offset: SupportsIndex = ...,
+        axis1: SupportsIndex = ...,
+        axis2: SupportsIndex = ...,
+        dtype: DTypeLike = ...,
+        out: None = ...,
+    ) -> Any: ...
+    @overload
+    def trace(
+        self,  # >= 2D MaskedArray
+        offset: SupportsIndex = ...,
+        axis1: SupportsIndex = ...,
+        axis2: SupportsIndex = ...,
+        dtype: DTypeLike = ...,
+        *,
+        out: _ArrayT,
+    ) -> _ArrayT: ...
+    @overload
+    def trace(
+        self,  # >= 2D MaskedArray
+        offset: SupportsIndex,
+        axis1: SupportsIndex,
+        axis2: SupportsIndex,
+        dtype: DTypeLike,
+        out: _ArrayT,
+    ) -> _ArrayT: ...
+
     def dot(self, b, out=..., strict=...): ...
     def sum(self, axis=..., dtype=..., out=..., keepdims=...): ...
-    def cumsum(self, axis=..., dtype=..., out=...): ...
+
+    @overload  # out: None (default)
+    def cumsum(self, /, axis: SupportsIndex | None = None, dtype: DTypeLike | None = None, out: None = None) -> _MaskedArray[Any]: ...
+    @overload  # out: ndarray
+    def cumsum(self, /, axis: SupportsIndex | None, dtype: DTypeLike | None, out: _ArrayT) -> _ArrayT: ...
+    @overload
+    def cumsum(self, /, axis: SupportsIndex | None = None, dtype: DTypeLike | None = None, *, out: _ArrayT) -> _ArrayT: ...
+
     def prod(self, axis=..., dtype=..., out=..., keepdims=...): ...
     product: Any
-    def cumprod(self, axis=..., dtype=..., out=...): ...
+
+    @overload  # out: None (default)
+    def cumprod(self, /, axis: SupportsIndex | None = None, dtype: DTypeLike | None = None, out: None = None) -> _MaskedArray[Any]: ...
+    @overload  # out: ndarray
+    def cumprod(self, /, axis: SupportsIndex | None, dtype: DTypeLike | None, out: _ArrayT) -> _ArrayT: ...
+    @overload
+    def cumprod(self, /, axis: SupportsIndex | None = None, dtype: DTypeLike | None = None, *, out: _ArrayT) -> _ArrayT: ...
+
     def mean(self, axis=..., dtype=..., out=..., keepdims=...): ...
     def anom(self, axis=..., dtype=...): ...
     def var(self, axis=..., dtype=..., out=..., ddof=..., keepdims=...): ...
     def std(self, axis=..., dtype=..., out=..., ddof=..., keepdims=...): ...
-    def round(self, decimals=..., out=...): ...
+
+    # Keep in sync with `ndarray.round`
+    @overload  # out=None (default)
+    def round(self, /, decimals: SupportsIndex = 0, out: None = None) -> Self: ...
+    @overload  # out=ndarray
+    def round(self, /, decimals: SupportsIndex, out: _ArrayT) -> _ArrayT: ...
+    @overload
+    def round(self, /, decimals: SupportsIndex = 0, *, out: _ArrayT) -> _ArrayT: ...
+
     def argsort(self, axis=..., kind=..., order=..., endwith=..., fill_value=..., *, stable=...): ...
 
     # Keep in-sync with np.ma.argmin

diff --git a/numpy/typing/tests/data/reveal/ma.pyi b/numpy/typing/tests/data/reveal/ma.pyi
@@ -383,6 +383,18 @@ assert_type(MAR_2d_f4.T, np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]
 assert_type(MAR_2d_f4.nonzero(), tuple[_Array1D[np.intp], ...])
 assert_type(MAR_2d_f4.nonzero()[0], _Array1D[np.intp])
 
+assert_type(MAR_f8.trace(), Any)
+assert_type(MAR_f8.trace(out=MAR_subclass), MaskedArraySubclass)
+
+assert_type(MAR_f8.round(), MaskedArray[np.float64])
+assert_type(MAR_f8.round(out=MAR_subclass), MaskedArraySubclass)
+
+assert_type(MAR_f8.cumprod(), MaskedArray[Any])
+assert_type(MAR_f8.cumprod(out=MAR_subclass), MaskedArraySubclass)
+
+assert_type(MAR_f8.cumsum(), MaskedArray[Any])
+assert_type(MAR_f8.cumsum(out=MAR_subclass), MaskedArraySubclass)
+
 # Masked Array addition
 
 assert_type(MAR_b + AR_LIKE_u, MaskedArray[np.uint32])