TYP: Type np.ma.is_mask (numpy#28833)

Author: MarcoGorelli

numpy/ma/core.pyi

@@ -5,7 +5,7 @@ from collections.abc import Sequence
 from typing import Any, Literal, SupportsIndex, TypeAlias, TypeVar, overload
 
 from _typeshed import Incomplete
-from typing_extensions import deprecated
+from typing_extensions import TypeIs, deprecated
 
 import numpy as np
 from numpy import (
@@ -341,7 +341,11 @@ def getmask(a): ...
 get_mask = getmask
 
 def getmaskarray(arr): ...
-def is_mask(m): ...
+
+# It's sufficient for `m` to have dtype with type: `type[np.bool_]`,
+# which isn't necessarily a ndarray. Please open an issue if this causes issues.
+def is_mask(m: object) -> TypeIs[NDArray[bool_]]: ...
+
 def make_mask(m, copy=..., shrink=..., dtype=...): ...
 def make_mask_none(newshape, dtype=...): ...
 def mask_or(m1, m2, copy=..., shrink=...): ...

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

@@ -10,6 +10,7 @@ MaskedNDArray: TypeAlias = np.ma.MaskedArray[_Shape, dtype[_ScalarT_co]]
 
 class MaskedNDArraySubclass(MaskedNDArray[np.complex128]): ...
 
+AR_b: NDArray[np.bool]
 AR_f4: NDArray[np.float32]
 AR_dt64: NDArray[np.datetime64]
 AR_td64: NDArray[np.timedelta64]
@@ -281,6 +282,15 @@ assert_type(np.ma.allclose(AR_f4, MAR_f4), bool)
 assert_type(np.ma.allclose(AR_f4, MAR_f4, masked_equal=False), bool)
 assert_type(np.ma.allclose(AR_f4, MAR_f4, rtol=.4, atol=.3), bool)
 
+assert_type(np.ma.is_mask(MAR_1d), bool)
+assert_type(np.ma.is_mask(AR_b), bool)
+
+def func(x: object) -> None:
+    if np.ma.is_mask(x):
+        assert_type(x, NDArray[np.bool])
+    else:
+        assert_type(x, object)
+
 assert_type(np.ma.nomask, np.bool[Literal[False]])
 # https://github.com/python/mypy/issues/18974
 assert_type(np.ma.MaskType, type[np.bool])  # type: ignore[assert-type]