TYP: Type MaskedArray.compress (numpy#29480)

Author: MarcoGorelli

numpy/ma/core.pyi

@@ -614,7 +614,36 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     set_fill_value: Any
     def filled(self, /, fill_value: _ScalarLike_co | None = None) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
     def compressed(self) -> ndarray[tuple[int], _DTypeT_co]: ...
-    def compress(self, condition, axis=..., out=...): ...
+
+    @overload
+    def compress(
+        self,
+        condition: _ArrayLikeBool_co,
+        axis: _ShapeLike | None,
+        out: _ArrayT
+    ) -> _ArrayT: ...
+    @overload
+    def compress(
+        self,
+        condition: _ArrayLikeBool_co,
+        axis: _ShapeLike | None = None,
+        *,
+        out: _ArrayT
+    ) -> _ArrayT: ...
+    @overload
+    def compress(
+        self,
+        condition: _ArrayLikeBool_co,
+        axis: None = None,
+        out: None = None
+    ) -> MaskedArray[tuple[int], _DTypeT_co]: ...
+    @overload
+    def compress(
+        self,
+        condition: _ArrayLikeBool_co,
+        axis: _ShapeLike | None = None,
+        out: None = None
+    ) -> MaskedArray[_AnyShape, _DTypeT_co]: ...
 
     # TODO: How to deal with the non-commutative nature of `==` and `!=`?
     # xref numpy/numpy#17368

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

@@ -301,6 +301,11 @@ assert_type(np.ma.count(MAR_o, None, True), NDArray[np.int_])
 
 assert_type(MAR_f4.compressed(), np.ndarray[tuple[int], np.dtype[np.float32]])
 
+assert_type(MAR_f4.compress([True, False]), np.ma.MaskedArray[tuple[int], np.dtype[np.float32]])
+assert_type(MAR_f4.compress([True, False], axis=0), MaskedArray[np.float32])
+assert_type(MAR_f4.compress([True, False], axis=0, out=MAR_subclass), MaskedArraySubclassC)
+assert_type(MAR_f4.compress([True, False], 0, MAR_subclass), MaskedArraySubclassC)
+
 assert_type(np.ma.compressed(MAR_i8), np.ndarray[tuple[int], np.dtype[np.int64]])
 assert_type(np.ma.compressed([[1, 2, 3]]), np.ndarray[tuple[int], np.dtype])