TYP: Type np.ma.{is_masked,ndim,size,ids,is_contiguous} (numpy#28715)

MarcoGorelli · web-flow · commit 030989013d41 · 2025-04-15T18:25:02.000+02:00
diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi
@@ -457,8 +457,8 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     def reshape(self, *s, **kwargs): ...
     def resize(self, newshape, refcheck=..., order=...): ...
     def put(self, indices, values, mode=...): ...
-    def ids(self): ...
-    def iscontiguous(self): ...
+    def ids(self) -> tuple[int, int]: ...
+    def iscontiguous(self) -> bool: ...
     def all(self, axis=..., out=..., keepdims=...): ...
     def any(self, axis=..., out=..., keepdims=...): ...
     def nonzero(self): ...
@@ -804,7 +804,7 @@ def array(
     subok=...,
     ndmin=...,
 ): ...
-def is_masked(x): ...
+def is_masked(x: object) -> bool: ...
 
 class _extrema_operation(_MaskedUFunc):
     compare: Any
@@ -1113,9 +1113,9 @@ def putmask(a, mask, values): ...
 def transpose(a, axes=...): ...
 def reshape(a, new_shape, order=...): ...
 def resize(x, new_shape): ...
-def ndim(obj): ...
+def ndim(obj: ArrayLike) -> int: ...
 def shape(obj): ...
-def size(obj, axis=...): ...
+def size(obj: ArrayLike, axis: SupportsIndex | None = None) -> int: ...
 def diff(a, /, n=..., axis=..., prepend=..., append=...): ...
 def where(condition, x=..., y=...): ...
 def choose(indices, choices, out=..., mode=...): ...
diff --git a/numpy/typing/tests/data/fail/ma.pyi b/numpy/typing/tests/data/fail/ma.pyi
@@ -93,11 +93,14 @@ m.argpartition(kind='cabbage')  # E: No overload variant
 m.argpartition(order=lambda: 'cabbage')  # E: No overload variant
 m.argpartition(AR_b)  # E: No overload variant
 
+np.ma.ndim(lambda: 'lambda')  # E: No overload variant
+
+np.ma.size(AR_b, axis='0')  # E: No overload variant
+
 m >= (lambda x: 'mango') # E: No overload variant
 
 m > (lambda x: 'mango') # E: No overload variant
 
 m <= (lambda x: 'mango') # E: No overload variant
 
 m < (lambda x: 'mango') # E: No overload variant
-
diff --git a/numpy/typing/tests/data/reveal/ma.pyi b/numpy/typing/tests/data/reveal/ma.pyi
@@ -30,6 +30,7 @@ MAR_subclass: MaskedNDArraySubclass
 
 MAR_1d: np.ma.MaskedArray[tuple[int], np.dtype[Any]]
 
+b: np.bool
 f4: np.float32
 f: float
 
@@ -168,6 +169,20 @@ assert_type(MAR_f4.partition(1, axis=0, kind='introselect', order='K'), None)
 assert_type(MAR_f4.argpartition(1), MaskedNDArray[np.intp])
 assert_type(MAR_1d.argpartition(1, axis=0, kind='introselect', order='K'), MaskedNDArray[np.intp])
 
+assert_type(np.ma.ndim(f4), int)
+assert_type(np.ma.ndim(MAR_b), int)
+assert_type(np.ma.ndim(AR_f4), int)
+
+assert_type(np.ma.size(b), int)
+assert_type(np.ma.size(MAR_f4, axis=0), int)
+assert_type(np.ma.size(AR_f4), int)
+
+assert_type(np.ma.is_masked(MAR_f4), bool)
+
+assert_type(MAR_f4.ids(), tuple[int, int])
+
+assert_type(MAR_f4.iscontiguous(), bool)
+
 assert_type(MAR_f4 >= 3, MaskedNDArray[np.bool])
 assert_type(MAR_i8 >= AR_td64, MaskedNDArray[np.bool])
 assert_type(MAR_b >= AR_td64, MaskedNDArray[np.bool])
