ENH: add dict to return type of func argument of DataFrame#apply when result_type="expand"

Author: skatsuta

pandas-stubs/core/frame.pyi

@@ -1126,7 +1126,7 @@ class DataFrame(NDFrame, OpsMixin):
     @overload
     def apply(
         self,
-        f: Callable[..., ListLikeExceptSeriesAndStr | Series],
+        f: Callable[..., ListLikeExceptSeriesAndStr | Series | dict],
         axis: AxisType = ...,
         raw: _bool = ...,
         args=...,

tests/test_frame.py

@@ -477,6 +477,9 @@ def returns_listlike_of_2(x: pd.Series) -> tuple[int, int]:
     def returns_listlike_of_3(x: pd.Series) -> tuple[int, int, int]:
         return (7, 8, 9)
 
+    def returns_dict(x: pd.Series) -> dict[str, int]:
+        return {"col4": 7, "col5": 8}
+
     # Misc checks
     check(assert_type(df.apply(np.exp), pd.DataFrame), pd.DataFrame)
     check(assert_type(df.apply(str), "pd.Series[str]"), pd.Series, str)
@@ -518,6 +521,10 @@ def gethead(s: pd.Series, y: int) -> pd.Series:
         ),
         pd.DataFrame,
     )
+    check(
+        assert_type(df.apply(returns_dict, result_type="expand"), pd.DataFrame),
+        pd.DataFrame,
+    )
 
     # Check various return types for result_type="reduce" with default axis (0)
     check(
@@ -599,6 +606,10 @@ def gethead(s: pd.Series, y: int) -> pd.Series:
         ),
         pd.DataFrame,
     )
+    check(
+        assert_type(df.apply(returns_dict, axis=1, result_type="expand"), pd.DataFrame),
+        pd.DataFrame,
+    )
 
     # Check various return types for result_type="reduce" with axis=1
     check(