TYP: Type MaskedArray.reshape (#29404)

MarcoGorelli · jorenham · web-flow · commit 8b2321d03c1a · 2025-07-23T19:31:05.000+02:00
Co-authored-by: Joren Hammudoglu &lt;jhammudoglu@gmail.com&gt;
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -2468,6 +2468,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
     def flatten(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], _DTypeT_co]: ...
     def ravel(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], _DTypeT_co]: ...
 
+    # Keep in sync with `MaskedArray.reshape`
     # NOTE: reshape also accepts negative integers, so we can't use integer literals
     @overload  # (None)
     def reshape(self, shape: None, /, *, order: _OrderACF = "C", copy: builtins.bool | None = None) -> Self: ...
diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi
@@ -17,9 +17,11 @@ from typing_extensions import TypeIs, TypeVar
 
 import numpy as np
 from numpy import (
+    _AnyShapeT,
     _HasDType,
     _HasDTypeWithRealAndImag,
     _ModeKind,
+    _OrderACF,
     _OrderKACF,
     _PartitionKind,
     _SortKind,
@@ -1126,7 +1128,90 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     def count(self, axis: _ShapeLike | None, keepdims: Literal[True]) -> NDArray[int_]: ...
 
     def ravel(self, order: _OrderKACF = "C") -> MaskedArray[tuple[int], _DTypeT_co]: ...
-    def reshape(self, *s, **kwargs): ...
+
+    # Keep in sync with `ndarray.reshape`
+    # NOTE: reshape also accepts negative integers, so we can't use integer literals
+    @overload  # (None)
+    def reshape(self, shape: None, /, *, order: _OrderACF = "C", copy: bool | None = None) -> Self: ...
+    @overload  # (empty_sequence)
+    def reshape(  # type: ignore[overload-overlap]  # mypy false positive
+        self,
+        shape: Sequence[Never],
+        /,
+        *,
+        order: _OrderACF = "C",
+        copy: bool | None = None,
+    ) -> MaskedArray[tuple[()], _DTypeT_co]: ...
+    @overload  # (() | (int) | (int, int) | ....)  # up to 8-d
+    def reshape(
+        self,
+        shape: _AnyShapeT,
+        /,
+        *,
+        order: _OrderACF = "C",
+        copy: bool | None = None,
+    ) -> MaskedArray[_AnyShapeT, _DTypeT_co]: ...
+    @overload  # (index)
+    def reshape(
+        self,
+        size1: SupportsIndex,
+        /,
+        *,
+        order: _OrderACF = "C",
+        copy: bool | None = None,
+    ) -> MaskedArray[tuple[int], _DTypeT_co]: ...
+    @overload  # (index, index)
+    def reshape(
+        self,
+        size1: SupportsIndex,
+        size2: SupportsIndex,
+        /,
+        *,
+        order: _OrderACF = "C",
+        copy: bool | None = None,
+    ) -> MaskedArray[tuple[int, int], _DTypeT_co]: ...
+    @overload  # (index, index, index)
+    def reshape(
+        self,
+        size1: SupportsIndex,
+        size2: SupportsIndex,
+        size3: SupportsIndex,
+        /,
+        *,
+        order: _OrderACF = "C",
+        copy: bool | None = None,
+    ) -> MaskedArray[tuple[int, int, int], _DTypeT_co]: ...
+    @overload  # (index, index, index, index)
+    def reshape(
+        self,
+        size1: SupportsIndex,
+        size2: SupportsIndex,
+        size3: SupportsIndex,
+        size4: SupportsIndex,
+        /,
+        *,
+        order: _OrderACF = "C",
+        copy: bool | None = None,
+    ) -> MaskedArray[tuple[int, int, int, int], _DTypeT_co]: ...
+    @overload  # (int, *(index, ...))
+    def reshape(
+        self,
+        size0: SupportsIndex,
+        /,
+        *shape: SupportsIndex,
+        order: _OrderACF = "C",
+        copy: bool | None = None,
+    ) -> MaskedArray[_AnyShape, _DTypeT_co]: ...
+    @overload  # (sequence[index])
+    def reshape(
+        self,
+        shape: Sequence[SupportsIndex],
+        /,
+        *,
+        order: _OrderACF = "C",
+        copy: bool | None = None,
+    ) -> MaskedArray[_AnyShape, _DTypeT_co]: ...
+
     def resize(self, newshape: Never, refcheck: bool = True, order: bool = False) -> NoReturn: ...
     def put(self, indices: _ArrayLikeInt_co, values: ArrayLike, mode: _ModeKind = "raise") -> None: ...
     def ids(self) -> tuple[int, int]: ...
diff --git a/numpy/typing/tests/data/reveal/ma.pyi b/numpy/typing/tests/data/reveal/ma.pyi
@@ -27,6 +27,7 @@ AR_LIKE_dt64: list[np.datetime64]
 AR_LIKE_o: list[np.object_]
 AR_number: NDArray[np.number]
 
+MAR_c8: MaskedArray[np.complex64]
 MAR_c16: MaskedArray[np.complex128]
 MAR_b: MaskedArray[np.bool]
 MAR_f4: MaskedArray[np.float32]
@@ -399,6 +400,13 @@ assert_type(MAR_f8.trace(out=MAR_subclass, dtype=None), MaskedArraySubclass)
 assert_type(MAR_f8.round(), MaskedArray[np.float64])
 assert_type(MAR_f8.round(out=MAR_subclass), MaskedArraySubclass)
 
+assert_type(MAR_i8.reshape(None), MaskedArray[np.int64])
+assert_type(MAR_f8.reshape(-1), np.ma.MaskedArray[tuple[int], np.dtype[np.float64]])
+assert_type(MAR_c8.reshape(2, 3, 4, 5), np.ma.MaskedArray[tuple[int, int, int, int], np.dtype[np.complex64]])
+assert_type(MAR_td64.reshape(()), np.ma.MaskedArray[tuple[()], np.dtype[np.timedelta64]])
+assert_type(MAR_s.reshape([]), np.ma.MaskedArray[tuple[()], np.dtype[np.str_]])
+assert_type(MAR_V.reshape((480, 720, 4)), np.ma.MaskedArray[tuple[int, int, int], np.dtype[np.void]])
+
 assert_type(MAR_f8.cumprod(), MaskedArray[Any])
 assert_type(MAR_f8.cumprod(out=MAR_subclass), MaskedArraySubclass)
 
