TYP: Type masked array shape, dtype, __int__, and __float__

Author: MarcoGorelli

numpy/ma/core.pyi

@@ -190,7 +190,9 @@ __all__ = [
     "zeros_like",
 ]
 
+_ShapeType = TypeVar("_ShapeType", bound=tuple[int, ...])
 _ShapeType_co = TypeVar("_ShapeType_co", bound=tuple[int, ...], covariant=True)
+_DType = TypeVar("_DType", bound=dtype[Any])
 _DType_co = TypeVar("_DType_co", bound=dtype[Any], covariant=True)
 
 MaskType = bool
@@ -347,13 +349,13 @@ class MaskedArray(ndarray[_ShapeType_co, _DType_co]):
     def __getitem__(self, indx): ...
     def __setitem__(self, indx, value): ...
     @property
-    def dtype(self): ...
+    def dtype(self) -> _DType_co: ...
     @dtype.setter
-    def dtype(self, dtype): ...
+    def dtype(self: MaskedArray[_ShapeType_co, _DType], dtype: _DType) -> None: ...
     @property
-    def shape(self): ...
+    def shape(self) -> _ShapeType_co: ...
     @shape.setter
-    def shape(self, shape): ...
+    def shape(self: MaskedArray[_ShapeType, _DType_co], shape: _ShapeType) -> None: ...
     def __setmask__(self, mask, copy=...): ...
     @property
     def mask(self): ...
@@ -413,8 +415,8 @@ class MaskedArray(ndarray[_ShapeType_co, _DType_co]):
     def __ifloordiv__(self, other): ...
     def __itruediv__(self, other): ...
     def __ipow__(self, other): ...
-    def __float__(self): ...
-    def __int__(self): ...
+    def __float__(self) -> float: ...
+    def __int__(self) -> int: ...
     @property  # type: ignore[misc]
     def imag(self): ...
     get_imag: Any

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

@@ -0,0 +1,9 @@
+from typing import Any
+
+import numpy as np
+import numpy.ma
+
+m: np.ma.MaskedArray[tuple[int], np.dtype[np.float64]]
+
+m.shape = (3, 1)  # E: Incompatible types in assignment
+m.dtype = np.bool  # E: Incompatible types in assignment

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

@@ -0,0 +1,12 @@
+import numpy as np
+from typing_extensions import assert_type
+
+
+m: np.ma.MaskedArray[tuple[int], np.dtype[np.float64]]
+
+assert_type(m.shape, tuple[int])
+
+assert_type(m.dtype, np.dtype[np.float64])
+
+assert_type(int(m), int)
+assert_type(float(m), float)