Support Callable conditions for DataFrame.where() and Series.where() (#310)

Author: gandhis1

pandas-stubs/core/frame.pyi

@@ -2017,7 +2017,11 @@ class DataFrame(NDFrame, OpsMixin):
     ) -> Series: ...
     def where(
         self,
-        cond: Series | DataFrame | np.ndarray,
+        cond: Series
+        | DataFrame
+        | np.ndarray
+        | Callable[[DataFrame], DataFrame]
+        | Callable[[Any], _bool],
         other=...,
         inplace: _bool = ...,
         axis: AxisType | None = ...,

pandas-stubs/core/series.pyi

@@ -1087,7 +1087,11 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     ) -> Series: ...
     def where(
         self,
-        cond: Series[S1] | Series[_bool] | np.ndarray,
+        cond: Series[S1]
+        | Series[_bool]
+        | np.ndarray
+        | Callable[[Series[S1]], Series[bool]]
+        | Callable[[S1], bool],
         other=...,
         inplace: _bool = ...,
         axis: SeriesAxisType | None = ...,

tests/test_frame.py

@@ -1764,3 +1764,20 @@ def test_loc_slice() -> None:
         index=pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=["num", "let"]),
     )
     check(assert_type(df1.loc[1, :], pd.DataFrame), pd.DataFrame)
+
+
+def test_where() -> 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.where(cond1), pd.DataFrame), pd.DataFrame)
+
+    def cond2(x: pd.DataFrame) -> pd.DataFrame:
+        return x > 1
+
+    check(assert_type(df.where(cond2), pd.DataFrame), pd.DataFrame)
+
+    cond3 = pd.DataFrame({"a": [True, True, False], "b": [False, False, False]})
+    check(assert_type(df.where(cond3), pd.DataFrame), pd.DataFrame)

tests/test_series.py

@@ -1192,3 +1192,20 @@ def test_types_to_numpy() -> None:
     check(assert_type(s.to_numpy(), np.ndarray), np.ndarray)
     check(assert_type(s.to_numpy(dtype="str", copy=True), np.ndarray), np.ndarray)
     check(assert_type(s.to_numpy(na_value=0), np.ndarray), np.ndarray)
+
+
+def test_where() -> None:
+    s = pd.Series([1, 2, 3], dtype=int)
+
+    def cond1(x: int) -> bool:
+        return x % 2 == 0
+
+    check(assert_type(s.where(cond1, other=0), "pd.Series[int]"), pd.Series, int)
+
+    def cond2(x: pd.Series[int]) -> pd.Series[bool]:
+        return x > 1
+
+    check(assert_type(s.where(cond2, other=0), "pd.Series[int]"), pd.Series, int)
+
+    cond3 = pd.Series([False, True, True])
+    check(assert_type(s.where(cond3, other=0), "pd.Series[int]"), pd.Series, int)