allow args and kwargs in groupby.apply

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,
+ )