add Series.case_when (#957)

add case_when

Author: Dr-Irv

pandas-stubs/core/series.pyi

@@ -1426,6 +1426,17 @@ class Series(IndexOpsMixin[S1], NDFrame):
         axis: AxisIndex | None = ...,
         level: Level | None = ...,
     ) -> Series[S1]: ...
+    def case_when(
+        self,
+        caselist: list[
+            tuple[
+                Sequence[bool]
+                | Series[bool]
+                | Callable[[Series], Series | np.ndarray | Sequence[bool]],
+                ListLikeU | Scalar | Callable[[Series], Series | np.ndarray],
+            ],
+        ],
+    ) -> Series: ...
     def truncate(
         self,
         before: date | _str | int | None = ...,

tests/test_series.py

@@ -3288,3 +3288,16 @@ def callable(x: int | NAType) -> str | NAType:
 
     series = pd.Series(["a", "b", "c"])
     check(assert_type(s.map(series, na_action=None), "pd.Series[str]"), pd.Series, str)
+
+
+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)