TYP: Type MaskedArray.__deepcopy__ and MaskedArray.argsort

MarcoGorelli · MarcoGorelli · commit dc6b3c8a2b35 · 2025-07-22T17:27:01.000+01:00
diff --git a/numpy/ma/core.py b/numpy/ma/core.py
@@ -5628,7 +5628,7 @@ def argsort(self, axis=np._NoValue, kind=None, order=None, endwith=True,
             is used.
         kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}, optional
             The sorting algorithm used.
-        order : list, optional
+        order : str or list of str, optional
             When `a` is an array with fields defined, this argument specifies
             which fields to compare first, second, etc.  Not all fields need be
             specified.
diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi
@@ -1398,7 +1398,16 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     @overload
     def round(self, /, decimals: SupportsIndex = 0, *, out: _ArrayT) -> _ArrayT: ...
 
-    def argsort(self, axis=..., kind=..., order=..., endwith=..., fill_value=..., *, stable=...): ...
+    def argsort(
+        self,
+        axis: SupportsIndex | _NoValueType = ...,
+        kind: _SortKind | None = None,
+        order: str | Sequence[str] | None = None,
+        endwith: bool = True,
+        fill_value: _ScalarLike_co | None = None,
+        *,
+        stable: bool = ...,
+    ) -> _MaskedArray[Any]: ...
 
     # Keep in-sync with np.ma.argmin
     @overload  # type: ignore[override]
@@ -1703,7 +1712,7 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
 
     #
     def __reduce__(self): ...
-    def __deepcopy__(self, memo=...): ...
+    def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self: ...
 
     # Keep `dtype` at the bottom to avoid name conflicts with `np.dtype`
     @property
diff --git a/numpy/typing/tests/data/fail/ma.pyi b/numpy/typing/tests/data/fail/ma.pyi
@@ -141,3 +141,5 @@ MAR_c //= 2  # type: ignore[misc]
 MAR_td64 **= 2  # type: ignore[misc]
 
 MAR_1d_f8.swapaxes(axis1=1, axis2=0)  # type: ignore[call-arg]
+
+MAR_1d_f8.argsort(axis=(1,0))  # type: ignore[arg-type]
diff --git a/numpy/typing/tests/data/reveal/ma.pyi b/numpy/typing/tests/data/reveal/ma.pyi
@@ -405,6 +405,12 @@ 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)
 
+assert_type(MAR_f8.__deepcopy__(), MaskedArray[np.float64])
+
+assert_type(MAR_f8.argsort(), MaskedArray[Any])
+assert_type(MAR_f8.argsort(axis=0, kind='heap', order=('x', 'y')), MaskedArray[Any])
+assert_type(MAR_f8.argsort(endwith=True, fill_value=1.5, stable=False), MaskedArray[Any])
+
 def invalid_resize() -> None:
     assert_type(MAR_f8.resize((1,1)), NoReturn)  # type: ignore[arg-type]
 
