case_when

Author: cmp0xff

pandas-stubs/core/series.pyi

@@ -220,6 +220,7 @@ from pandas._typing import (
     WriteBuffer,
     _T_co,
     np_1darray,
+    np_1darray_bool,
     np_1darray_dt,
     np_1darray_int64,
     np_1darray_intp,
@@ -1678,11 +1679,12 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
     ) -> Series[S1]: ...
     def case_when(
         self,
-        caselist: list[
+        caselist: Sequence[
             tuple[
                 Sequence[bool]
+                | np_1darray_bool
                 | Series[bool]
-                | Callable[[Series], Series | np_ndarray | Sequence[bool]],
+                | Callable[[Series], Sequence[bool] | np_1darray_bool | Series[bool]],
                 ListLikeU | Scalar | Callable[[Series], Series | np_ndarray],
             ],
         ],

tests/series/test_series.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import (
+    Callable,
     Hashable,
     Iterable,
     Iterator,
@@ -60,6 +61,7 @@
     check,
     ensure_clean,
     np_1darray,
+    np_1darray_bool,
     np_ndarray_num,
     pytest_warns_bounded,
 )
@@ -3623,13 +3625,20 @@ def test_case_when() -> None:
     c = pd.Series([6, 7, 8, 9], name="c")
     a = pd.Series([0, 0, 1, 2])
     b = pd.Series([0, 3, 4, 5])
-    r = c.case_when(
-        caselist=[
-            (a.gt(0), a),
-            (b.gt(0), b),
-        ]
-    )
-    check(assert_type(r, pd.Series), pd.Series)
+
+    c0 = [(a.gt(0), a), (b.gt(0), b)]
+    check(assert_type(c.case_when(c0), pd.Series), pd.Series)
+
+    def foo_factory(
+        thresh: int,
+    ) -> Callable[[pd.Series], pd.Series[bool] | np_1darray_bool]:
+        def foo(s: pd.Series) -> pd.Series[bool] | np_1darray_bool:
+            return s >= thresh
+
+        return foo
+
+    c1 = [(foo_factory(2), a), (foo_factory(0), b)]
+    check(assert_type(c.case_when(c1), pd.Series), pd.Series)
 
 
 def test_series_unique_timestamp() -> None: