allow args and kwargs in groupby.apply (#1268)

* allow args and kwargs in groupby.apply * use paramspec on args, kwargs * remove include_groups. use P.args and P.kwargs in apply defs * add test to see if failures are picked up

Author: Dr-Irv

pandas-stubs/core/groupby/generic.pyi

@@ -11,6 +11,7 @@ from typing import (
  Generic,
  Literal,
  NamedTuple,
+ Protocol,
  TypeVar,
  final,
  overload,
@@ -208,28 +209,45 @@ class SeriesGroupBy(GroupBy[Series[S2]], Generic[S2, ByT]):
 
 _TT = TypeVar("_TT", bound=Literal[True, False])
 
+# ty ignore needed because of https://github.com/astral-sh/ty/issues/157#issuecomment-3017337945
+class DFCallable1(Protocol[P]): # ty: ignore[invalid-argument-type]
+ def __call__(
+ self, df: DataFrame, /, *args: P.args, **kwargs: P.kwargs
+ ) -> Scalar | list | dict: ...
+
+class DFCallable2(Protocol[P]): # ty: ignore[invalid-argument-type]
+ def __call__(
+ self, df: DataFrame, /, *args: P.args, **kwargs: P.kwargs
+ ) -> DataFrame | Series: ...
+
+class DFCallable3(Protocol[P]): # ty: ignore[invalid-argument-type]
+ def __call__(self, df: Iterable, /, *args: P.args, **kwargs: P.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],
- *args,
- **kwargs,
+ func: DFCallable1[P],
+ /,
+ *args: P.args,
+ **kwargs: P.kwargs,
  ) -> Series: ...
  @overload
  def apply(
  self,
- func: Callable[[DataFrame], Series | DataFrame],
- *args,
- **kwargs,
+ func: DFCallable2[P],
+ /,
+ *args: P.args,
+ **kwargs: P.kwargs,
  ) -> DataFrame: ...
  @overload
- def apply( # pyright: ignore[reportOverlappingOverload]
+ def apply(
  self,
- func: Callable[[Iterable], float],
- *args,
- **kwargs,
+ func: DFCallable3[P],
+ /,
+ *args: P.args,
+ **kwargs: P.kwargs,
  ) -> DataFrame: ...
  # error: overload 1 overlaps overload 2 because of different return types
  @overload

tests/test_groupby.py

@@ -271,7 +271,7 @@ def resample_interpolate(x: DataFrame) -> DataFrame:
 
  check(
  assert_type(
- GB_DF.apply(resample_interpolate, include_groups=False),
+ GB_DF.apply(resample_interpolate),
  DataFrame,
  ),
  DataFrame,
@@ -284,7 +284,6 @@ def resample_interpolate_linear(x: DataFrame) -> DataFrame:
  assert_type(
  GB_DF.apply(
  resample_interpolate_linear,
- include_groups=False,
  ),
  DataFrame,
  ),
@@ -1099,3 +1098,28 @@ 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,
+ )
+ if TYPE_CHECKING_INVALID_USAGE:
+ df.groupby("group", group_keys=False)[["group", "value"]].apply(
+ add_constant_to_mean,
+ constant="5", # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType]
+ )