More detailed typing for DataFrame.apply.

Whether it returns a Series or a DataFrame depends on the return type of
the callable. In the case of the callable returning a scalar, the result
is a Series unless the result_type is "broadcast".

Author: danielroseman

pandas-stubs/core/frame.pyi

@@ -1088,14 +1088,32 @@ class DataFrame(NDFrame, OpsMixin):
         **kwargs,
     ) -> DataFrame: ...
     @overload
-    def apply(self, f: Callable) -> Series: ...
+    def apply(
+        self,
+        f: Callable[..., Series],
+        axis: AxisType = ...,
+        raw: _bool = ...,
+        result_type: Literal["expand", "reduce", "broadcast"] | None = ...,
+        args=...,
+        **kwargs,
+    ) -> DataFrame: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., Scalar],
+        axis: AxisType = ...,
+        raw: _bool = ...,
+        result_type: Literal["expand", "reduce"] | None = ...,
+        args=...,
+        **kwargs,
+    ) -> Series: ...
     @overload
     def apply(
         self,
-        f: Callable,
-        axis: AxisType,
+        f: Callable[..., Scalar],
+        result_type: Literal["broadcast"],
+        axis: AxisType = ...,
         raw: _bool = ...,
-        result_type: _str | None = ...,
         args=...,
         **kwargs,
     ) -> DataFrame: ...

tests/test_frame.py

@@ -460,9 +460,22 @@ def test_types_unique() -> None:
 
 def test_types_apply() -> None:
     df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
-    df.apply(lambda x: x**2)
-    df.apply(np.exp)
-    df.apply(str)
+
+    def returns_series(x: pd.Series) -> pd.Series:
+        return x**2
+
+    check(assert_type(df.apply(returns_series), pd.DataFrame), pd.DataFrame)
+
+    def returns_scalar(x: pd.Series) -> float:
+        return 2
+
+    check(assert_type(df.apply(returns_scalar), pd.Series), pd.Series)
+    check(
+        assert_type(df.apply(returns_scalar, result_type="broadcast"), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(assert_type(df.apply(np.exp), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.apply(str), pd.Series), pd.Series)
 
 
 def test_types_applymap() -> None: