improve typing and add test for callable

Author: samukweku

pandas/core/series.py

@@ -5638,7 +5638,12 @@ def between(
 
     def case_when(
         self,
-        caselist: list[tuple[ArrayLike, ArrayLike | Scalar]],
+        caselist: list[
+            tuple[
+                ArrayLike | Callable[[Series | np.ndarray | Sequence[bool]]],
+                ArrayLike | Callable[[Series | np.ndarray]] | Scalar,
+            ]
+        ],
     ) -> Series:
         """
         Replace values where the conditions are True.

pandas/tests/series/methods/test_case_when.py

@@ -129,3 +129,20 @@ def test_case_when_non_range_index():
     result = Series(5, index=df.index, name="A").case_when([(df.A.gt(0), df.B)])
     expected = df.A.mask(df.A.gt(0), df.B).where(df.A.gt(0), 5)
     tm.assert_series_equal(result, expected)
+
+
+def test_case_when_callable():
+    """
+    Test output on a callable
+    """
+    # https://numpy.org/doc/stable/reference/generated/numpy.piecewise.html
+    x = np.linspace(-2.5, 2.5, 6)
+    ser = Series(x)
+    result = ser.case_when(
+        caselist=[
+            (lambda df: df < 0, lambda df: -df),
+            (lambda df: df >= 0, lambda df: df),
+        ]
+    )
+    expected = np.piecewise(x, [x < 0, x >= 0], [lambda x: -x, lambda x: x])
+    tm.assert_series_equal(result, Series(expected))