TYP: Type MaskedArray.repeat, improve overloads for NDArray.repeat, generic.repeat, and np.repeat (numpy#28849)

Author: MarcoGorelli

numpy/__init__.pyi

@@ -2421,10 +2421,17 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
         mode: _ModeKind = ...,
     ) -> _ArrayT: ...
 
+    @overload
     def repeat(
         self,
         repeats: _ArrayLikeInt_co,
-        axis: SupportsIndex | None = ...,
+        axis: None = None,
+    ) -> ndarray[tuple[int], _DTypeT_co]: ...
+    @overload
+    def repeat(
+        self,
+        repeats: _ArrayLikeInt_co,
+        axis: SupportsIndex,
     ) -> ndarray[_Shape, _DTypeT_co]: ...
 
     def flatten(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], _DTypeT_co]: ...
@@ -3685,7 +3692,7 @@ class generic(_ArrayOrScalarCommon, Generic[_ItemT_co]):
         mode: _ModeKind = ...,
     ) -> _ArrayT: ...
 
-    def repeat(self, repeats: _ArrayLikeInt_co, axis: SupportsIndex | None = ...) -> NDArray[Self]: ...
+    def repeat(self, repeats: _ArrayLikeInt_co, axis: SupportsIndex | None = None) -> ndarray[tuple[int], dtype[Self]]: ...
     def flatten(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], dtype[Self]]: ...
     def ravel(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], dtype[Self]]: ...
 

numpy/_core/fromnumeric.pyi

@@ -277,13 +277,25 @@ def choose(
 def repeat(
     a: _ArrayLike[_ScalarT],
     repeats: _ArrayLikeInt_co,
-    axis: SupportsIndex | None = ...,
+    axis: None = None,
+) -> np.ndarray[tuple[int], np.dtype[_ScalarT]]: ...
+@overload
+def repeat(
+    a: _ArrayLike[_ScalarT],
+    repeats: _ArrayLikeInt_co,
+    axis: SupportsIndex,
 ) -> NDArray[_ScalarT]: ...
 @overload
 def repeat(
     a: ArrayLike,
     repeats: _ArrayLikeInt_co,
-    axis: SupportsIndex | None = ...,
+    axis: None = None,
+) -> np.ndarray[tuple[int], np.dtype[Any]]: ...
+@overload
+def repeat(
+    a: ArrayLike,
+    repeats: _ArrayLikeInt_co,
+    axis: SupportsIndex,
 ) -> NDArray[Any]: ...
 
 def put(

numpy/ma/core.pyi

@@ -778,7 +778,20 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     copy: Any
     diagonal: Any
     flatten: Any
-    repeat: Any
+
+    @overload
+    def repeat(
+        self,
+        repeats: _ArrayLikeInt_co,
+        axis: None = None,
+    ) -> MaskedArray[tuple[int], _DTypeT_co]: ...
+    @overload
+    def repeat(
+        self,
+        repeats: _ArrayLikeInt_co,
+        axis: SupportsIndex,
+    ) -> MaskedArray[_Shape, _DTypeT_co]: ...
+
     squeeze: Any
     swapaxes: Any
     T: Any

numpy/typing/tests/data/reveal/fromnumeric.pyi

@@ -46,11 +46,13 @@ assert_type(np.choose([1], [True, True]), npt.NDArray[Any])
 assert_type(np.choose([1], AR_b), npt.NDArray[np.bool])
 assert_type(np.choose([1], AR_b, out=AR_f4), npt.NDArray[np.float32])
 
-assert_type(np.repeat(b, 1), npt.NDArray[np.bool])
-assert_type(np.repeat(f4, 1), npt.NDArray[np.float32])
-assert_type(np.repeat(f, 1), npt.NDArray[Any])
-assert_type(np.repeat(AR_b, 1), npt.NDArray[np.bool])
-assert_type(np.repeat(AR_f4, 1), npt.NDArray[np.float32])
+assert_type(np.repeat(b, 1), np.ndarray[tuple[int], np.dtype[np.bool]])
+assert_type(np.repeat(b, 1, axis=0), npt.NDArray[np.bool])
+assert_type(np.repeat(f4, 1), np.ndarray[tuple[int], np.dtype[np.float32]])
+assert_type(np.repeat(f, 1), np.ndarray[tuple[int], np.dtype[Any]])
+assert_type(np.repeat(AR_b, 1), np.ndarray[tuple[int], np.dtype[np.bool]])
+assert_type(np.repeat(AR_f4, 1), np.ndarray[tuple[int], np.dtype[np.float32]])
+assert_type(np.repeat(AR_f4, 1, axis=0), npt.NDArray[np.float32])
 
 # TODO: array_bdd tests for np.put()
 

numpy/typing/tests/data/reveal/ma.pyi

@@ -276,6 +276,11 @@ assert_type(np.ma.filled([[1,2,3]]), NDArray[Any])
 # https://github.com/numpy/numpy/pull/28742#discussion_r2048968375
 assert_type(np.ma.filled(MAR_1d), np.ndarray[tuple[int], np.dtype])  # type: ignore[assert-type]
 
+assert_type(MAR_b.repeat(3), np.ma.MaskedArray[tuple[int], np.dtype[np.bool]])
+assert_type(MAR_2d_f4.repeat(MAR_i8), np.ma.MaskedArray[tuple[int], np.dtype[np.float32]])
+assert_type(MAR_2d_f4.repeat(MAR_i8, axis=None), np.ma.MaskedArray[tuple[int], np.dtype[np.float32]])
+assert_type(MAR_2d_f4.repeat(MAR_i8, axis=0), MaskedNDArray[np.float32])
+
 assert_type(np.ma.allequal(AR_f4, MAR_f4), bool)
 assert_type(np.ma.allequal(AR_f4, MAR_f4, fill_value=False), bool)
 

numpy/typing/tests/data/reveal/ndarray_misc.pyi

@@ -126,9 +126,12 @@ assert_type(f8.round(), np.float64)
 assert_type(AR_f8.round(), npt.NDArray[np.float64])
 assert_type(AR_f8.round(out=B), SubClass)
 
-assert_type(f8.repeat(1), npt.NDArray[np.float64])
-assert_type(AR_f8.repeat(1), npt.NDArray[np.float64])
-assert_type(B.repeat(1), npt.NDArray[np.object_])
+assert_type(f8.repeat(1), np.ndarray[tuple[int], np.dtype[np.float64]])
+assert_type(f8.repeat(1, axis=0), np.ndarray[tuple[int], np.dtype[np.float64]])
+assert_type(AR_f8.repeat(1), np.ndarray[tuple[int], np.dtype[np.float64]])
+assert_type(AR_f8.repeat(1, axis=0), npt.NDArray[np.float64])
+assert_type(B.repeat(1), np.ndarray[tuple[int], np.dtype[np.object_]])
+assert_type(B.repeat(1, axis=0), npt.NDArray[np.object_])
 
 assert_type(f8.std(), Any)
 assert_type(AR_f8.std(), Any)