allow args and kwargs in groupby.apply

Author: Dr-Irv

pandas-stubs/core/groupby/generic.pyi

@@ -11,6 +11,7 @@ from typing import (
     Generic,
     Literal,
     NamedTuple,
+    Protocol,
     TypeVar,
     final,
     overload,
@@ -208,26 +209,35 @@ class SeriesGroupBy(GroupBy[Series[S2]], Generic[S2, ByT]):
 
 _TT = TypeVar("_TT", bound=Literal[True, False])
 
+class DFCallable1(Protocol):
+    def __call__(self, df: DataFrame, /, *args, **kwargs) -> Scalar | list | dict: ...
+
+class DFCallable2(Protocol):
+    def __call__(self, df: DataFrame, /, *args, **kwargs) -> DataFrame | Series: ...
+
+class DFCallable3(Protocol):
+    def __call__(self, df: Iterable, /, *args, **kwargs) -> float: ...
+
 class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
     # error: Overload 3 for "apply" will never be used because its parameters overlap overload 1
     @overload  # type: ignore[override]
     def apply(
         self,
-        func: Callable[[DataFrame], Scalar | list | dict],
+        func: DFCallable1,
         *args,
         **kwargs,
     ) -> Series: ...
     @overload
     def apply(
         self,
-        func: Callable[[DataFrame], Series | DataFrame],
+        func: DFCallable2,
         *args,
         **kwargs,
     ) -> DataFrame: ...
     @overload
     def apply(  # pyright: ignore[reportOverlappingOverload]
         self,
-        func: Callable[[Iterable], float],
+        func: DFCallable3,
         *args,
         **kwargs,
     ) -> DataFrame: ...

tests/test_groupby.py

@@ -1102,3 +1102,23 @@ def test_dataframe_value_counts() -> None:
         Series,
         np.int64,
     )
+
+
+def test_dataframe_apply_kwargs() -> None:
+    # GH 1266
+    df = DataFrame({"group": ["A", "A", "B", "B", "C"], "value": [10, 15, 10, 25, 30]})
+
+    def add_constant_to_mean(group: DataFrame, constant: int) -> DataFrame:
+        mean_val = group["value"].mean()
+        group["adjusted"] = mean_val + constant
+        return group
+
+    check(
+        assert_type(
+            df.groupby("group", group_keys=False)[["group", "value"]].apply(
+                add_constant_to_mean, constant=5
+            ),
+            DataFrame,
+        ),
+        DataFrame,
+    )