Add callable condition type in Series.mask, DataFrame.mask (#918)

* Add callable condition type in Series.mask, DataFrame.mask

* Get rid of ignored rule (pyright)

Author: georgesittas

pandas-stubs/core/frame.pyi

@@ -1833,7 +1833,13 @@ class DataFrame(NDFrame, OpsMixin):
     def lt(self, other, axis: Axis = ..., level: Level | None = ...) -> DataFrame: ...
     def mask(
         self,
-        cond: Series | DataFrame | np.ndarray,
+        cond: (
+            Series
+            | DataFrame
+            | np.ndarray
+            | Callable[[DataFrame], DataFrame]
+            | Callable[[Any], _bool]
+        ),
         other=...,
         *,
         inplace: _bool = ...,

pandas-stubs/core/groupby/generic.pyi

@@ -199,7 +199,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT]):
         **kwargs,
     ) -> DataFrame: ...
     @overload
-    def apply(  # pyright: ignore[reportOverlappingOverload,reportIncompatibleMethodOverride]
+    def apply(  # pyright: ignore[reportOverlappingOverload]
         self,
         func: Callable[[Iterable], float],
         *args,

pandas-stubs/core/series.pyi

@@ -1433,7 +1433,13 @@ class Series(IndexOpsMixin[S1], NDFrame):
     ) -> Series[S1]: ...
     def mask(
         self,
-        cond: MaskType,
+        cond: (
+            Series[S1]
+            | Series[_bool]
+            | np.ndarray
+            | Callable[[Series[S1]], Series[bool]]
+            | Callable[[S1], bool]
+        ),
         other: Scalar | Series[S1] | DataFrame | Callable | NAType | None = ...,
         *,
         inplace: _bool = ...,

tests/test_frame.py

@@ -2755,6 +2755,15 @@ def cond2(x: pd.DataFrame) -> pd.DataFrame:
     check(assert_type(df.where(cond3), pd.DataFrame), pd.DataFrame)
 
 
+def test_mask() -> None:
+    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+
+    def cond1(x: int) -> bool:
+        return x % 2 == 0
+
+    check(assert_type(df.mask(cond1), pd.DataFrame), pd.DataFrame)
+
+
 def test_setitem_loc() -> None:
     # GH 254
     df = pd.DataFrame.from_dict(

tests/test_series.py

@@ -2823,6 +2823,12 @@ def test_types_mask() -> None:
     # Test case with a boolean condition and a scalar value
     check(assert_type(s.mask(s > 3, 10), "pd.Series[int]"), pd.Series, np.integer)
 
+    def cond(x: int) -> bool:
+        return x % 2 == 0
+
+    # Test case with a callable condition and a scalar value
+    check(assert_type(s.mask(cond, 10), "pd.Series[int]"), pd.Series, np.integer)
+
     # Test case with a boolean condition and a callable
     def double(x):
         return x * 2