Repository: pandas-dev/pandas-stubs (language: python)

Issue:
Passing a callable to DataFrame.mask's 'cond' argument results in arg-type error (mypy)
**Describe the bug**

The `cond` argument of the `DataFrame.mask` method currently [doesn't expect a callable](https://github.com/pandas-dev/pandas-stubs/blob/e63d1cbf5087cd7cd934db1c84184cb2caacbac9/pandas-stubs/core/frame.pyi#L1836), even though it should, according to the method's [documentation page](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.mask.html):

> If cond is callable, it is computed on the Series/DataFrame and should return boolean Series/DataFrame or array. The callable must not change input Series/DataFrame (though pandas doesn’t check it).

**To Reproduce**

1. Minimally runnable `pandas` example that is not properly checked by the stubs:

```python
import pandas as pd
df = pd.DataFrame({"a": [1]}).mask(cond=lambda x: x.notna(), other=2)
print(df)
#    a
# 0  2
```

2. Type checker: `mypy`

3. Show the error message received from that type checker while checking your example.

```
(.venv) ➜  python -m mypy test.py
test.py:2: error: Argument "cond" to "mask" of "DataFrame" has incompatible type "Callable[[Any], Any]"; expected "Series[Any] | DataFrame | ndarray[Any, Any]"  [arg-type]
Found 1 error in 1 file (checked 1 source file)
```

**Please complete the following information:**
 - OS: MacOS
 - OS Version: Sonoma 14.2.1
 - python version: Python 3.10.14
 - version of type checker: 1.10.0
 - version of installed `pandas-stubs`: 2.2.1.240316

**Additional context**

I've also tried using `typing.cast` to cast the lambda into `Callable[[DataFrame], DataFrame]`, but it doesn't work, as expected. Happy to provide more details if needed.


Discussion / hints:
Thanks for the report.  It appears we have `Callable` in `Series.mask()`, but not `DataFrame.mask()`.

PR with tests welcome.

Gotcha! Appreciate the quick response, will try to make a PR for this 👍

Fix patch:
diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
index 78384247..f2deecaa 100644
--- a/pandas-stubs/core/frame.pyi
+++ b/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 = ...,
diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi
index 715899fc..37f3c0dc 100644
--- a/pandas-stubs/core/groupby/generic.pyi
+++ b/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,
diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
index 1d48feba..b6111db8 100644
--- a/pandas-stubs/core/series.pyi
+++ b/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 = ...,


Test patch:
diff --git a/tests/test_frame.py b/tests/test_frame.py
index f156a8e3..ed42dad8 100644
--- a/tests/test_frame.py
+++ b/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(
diff --git a/tests/test_series.py b/tests/test_series.py
index 203f8170..d2a8af78 100644
--- a/tests/test_series.py
+++ b/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