TYP: Type MaskedArray.count and np.ma.count (numpy#28735)

MarcoGorelli · web-flow · commit c75757c8236f · 2025-04-16T23:52:50.000+02:00
diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi
@@ -19,6 +19,7 @@ from numpy import (
     expand_dims,
     float64,
     generic,
+    int_,
     intp,
     ndarray,
 )
@@ -452,7 +453,17 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     @property  # type: ignore[misc]
     def real(self): ...
     get_real: Any
-    def count(self, axis=..., keepdims=...): ...
+
+    # keep in sync with `np.ma.count`
+    @overload
+    def count(self, axis: None = None, keepdims: Literal[False] | _NoValueType = ...) -> int: ...
+    @overload
+    def count(self, axis: _ShapeLike, keepdims: bool | _NoValueType = ...) -> NDArray[int_]: ...
+    @overload
+    def count(self, axis: _ShapeLike | None = ..., *, keepdims: Literal[True]) -> NDArray[int_]: ...
+    @overload
+    def count(self, axis: _ShapeLike | None, keepdims: Literal[True]) -> NDArray[int_]: ...
+
     def ravel(self, order=...): ...
     def reshape(self, *s, **kwargs): ...
     def resize(self, newshape, refcheck=..., order=...): ...
@@ -949,7 +960,15 @@ sum: _frommethod
 swapaxes: _frommethod
 trace: _frommethod
 var: _frommethod
-count: _frommethod
+
+@overload
+def count(self: ArrayLike, axis: None = None, keepdims: Literal[False] | _NoValueType = ...) -> int: ...
+@overload
+def count(self: ArrayLike, axis: _ShapeLike, keepdims: bool | _NoValueType = ...) -> NDArray[int_]: ...
+@overload
+def count(self: ArrayLike, axis: _ShapeLike | None = ..., *, keepdims: Literal[True]) -> NDArray[int_]: ...
+@overload
+def count(self: ArrayLike, axis: _ShapeLike | None, keepdims: Literal[True]) -> NDArray[int_]: ...
 
 @overload
 def argmin(
diff --git a/numpy/typing/tests/data/fail/ma.pyi b/numpy/typing/tests/data/fail/ma.pyi
@@ -104,3 +104,7 @@ 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.count(axis=0.)  # E: No overload variant
+
+np.ma.count(m, axis=0.)  # E: No overload variant
diff --git a/numpy/typing/tests/data/reveal/ma.pyi b/numpy/typing/tests/data/reveal/ma.pyi
@@ -232,3 +232,18 @@ assert_type(MAR_1d <= 0, MaskedNDArray[np.bool])
 assert_type(MAR_s <= MAR_s, MaskedNDArray[np.bool])
 assert_type(MAR_byte <= MAR_byte, MaskedNDArray[np.bool])
 
+assert_type(MAR_byte.count(), int)
+assert_type(MAR_f4.count(axis=None), int)
+assert_type(MAR_f4.count(axis=0), NDArray[np.int_])
+assert_type(MAR_b.count(axis=(0,1)), NDArray[np.int_])
+assert_type(MAR_o.count(keepdims=True), NDArray[np.int_])
+assert_type(MAR_o.count(axis=None, keepdims=True), NDArray[np.int_])
+assert_type(MAR_o.count(None, True), NDArray[np.int_])
+
+assert_type(np.ma.count(MAR_byte), int)
+assert_type(np.ma.count(MAR_byte, axis=None), int)
+assert_type(np.ma.count(MAR_f4, axis=0), NDArray[np.int_])
+assert_type(np.ma.count(MAR_b, axis=(0,1)), NDArray[np.int_])
+assert_type(np.ma.count(MAR_o, keepdims=True), NDArray[np.int_])
+assert_type(np.ma.count(MAR_o, axis=None, keepdims=True), NDArray[np.int_])
+assert_type(np.ma.count(MAR_o, None, True), NDArray[np.int_])
