TYP: Type MaskedArray.filled and np.ma.filled (numpy#28742)

MarcoGorelli · VascoConceicao · commit 7c1ba937092a · 2025-05-06T13:27:39.000+01:00
diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi
@@ -224,6 +224,7 @@ _DTypeT = TypeVar("_DTypeT", bound=dtype[Any])
 _DTypeT_co = TypeVar("_DTypeT_co", bound=dtype[Any], covariant=True)
 _ArrayT = TypeVar("_ArrayT", bound=ndarray[Any, Any])
 _ScalarT = TypeVar("_ScalarT", bound=generic)
+_ScalarT_co = TypeVar("_ScalarT_co", bound=generic)
 # A subset of `MaskedArray` that can be parametrized w.r.t. `np.generic`
 _MaskedArray: TypeAlias = MaskedArray[_Shape, dtype[_ScalarT]]
 
@@ -239,7 +240,12 @@ def minimum_fill_value(obj): ...
 def maximum_fill_value(obj): ...
 def set_fill_value(a, fill_value): ...
 def common_fill_value(a, b): ...
-def filled(a, fill_value=...): ...
+@overload
+def filled(a: ndarray[_ShapeT_co, _DTypeT_co], fill_value: _ScalarLike_co | None = None) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
+@overload
+def filled(a: _ArrayLike[_ScalarT_co], fill_value: _ScalarLike_co | None = None) -> NDArray[_ScalarT_co]: ...
+@overload
+def filled(a: ArrayLike, fill_value: _ScalarLike_co | None = None) -> NDArray[Any]: ...
 def getdata(a, subok=...): ...
 get_data = getdata
 
@@ -418,7 +424,7 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
     def fill_value(self, value=...): ...
     get_fill_value: Any
     set_fill_value: Any
-    def filled(self, fill_value=...): ...
+    def filled(self, /, fill_value: _ScalarLike_co | None = None) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
     def compressed(self): ...
     def compress(self, condition, axis=..., out=...): ...
     def __eq__(self, other): ...
diff --git a/numpy/typing/tests/data/reveal/ma.pyi b/numpy/typing/tests/data/reveal/ma.pyi
@@ -247,3 +247,13 @@ 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_])
+
+assert_type(MAR_f4.filled(float('nan')), NDArray[np.float32])
+assert_type(MAR_i8.filled(), NDArray[np.int64])
+assert_type(MAR_1d.filled(), np.ndarray[tuple[int], np.dtype[Any]])
+
+assert_type(np.ma.filled(MAR_f4, float('nan')), NDArray[np.float32])
+assert_type(np.ma.filled([[1,2,3]]), NDArray[Any])
+# PyRight detects this one correctly, but mypy doesn't.
+# https://github.com/numpy/numpy/pull/28742#discussion_r2048968375
+assert_type(np.ma.filled(MAR_1d), np.ndarray[tuple[int], np.dtype[Any]])  # type: ignore[assert-type]
