GH456 Attempt GroupBy.aggregate improved typing

Author: loicdiridollou

pandas-stubs/_typing.pyi

@@ -925,44 +925,6 @@ GroupByObjectNonScalar: TypeAlias = (
     | list[Grouper]
 )
 GroupByObject: TypeAlias = Scalar | Index | GroupByObjectNonScalar | Series
-GroupByFuncStrs: TypeAlias = Literal[
-    # Reduction/aggregation functions
-    "all",
-    "any",
-    "corrwith",
-    "count",
-    "first",
-    "idxmax",
-    "idxmin",
-    "last",
-    "max",
-    "mean",
-    "median",
-    "min",
-    "nunique",
-    "prod",
-    "quantile",
-    "sem",
-    "size",
-    "skew",
-    "std",
-    "sum",
-    "var",
-    # Transformation functions
-    "bfill",
-    "cumcount",
-    "cummax",
-    "cummin",
-    "cumprod",
-    "cumsum",
-    "diff",
-    "ffill",
-    "fillna",
-    "ngroup",
-    "pct_change",
-    "rank",
-    "shift",
-]
 
 StataDateFormat: TypeAlias = Literal[
     "tc",

pandas-stubs/core/groupby/base.pyi

@@ -1,7 +1,56 @@
 from collections.abc import Hashable
 import dataclasses
+from typing import (
+    Literal,
+    TypeAlias,
+)
 
 @dataclasses.dataclass(order=True, frozen=True)
 class OutputKey:
     label: Hashable
     position: int
+
+reduction_kernels: TypeAlias = Literal[
+    "all",
+    "any",
+    "corrwith",
+    "count",
+    "first",
+    "idxmax",
+    "idxmin",
+    "last",
+    "max",
+    "mean",
+    "median",
+    "min",
+    "nunique",
+    "prod",
+    # as long as `quantile`'s signature accepts only
+    # a single quantile value, it's a reduction.
+    # GH#27526 might change that.
+    "quantile",
+    "sem",
+    "size",
+    "skew",
+    "std",
+    "sum",
+    "var",
+]
+
+transformation_kernels: TypeAlias = Literal[
+    "bfill",
+    "cumcount",
+    "cummax",
+    "cummin",
+    "cumprod",
+    "cumsum",
+    "diff",
+    "ffill",
+    "fillna",
+    "ngroup",
+    "pct_change",
+    "rank",
+    "shift",
+]
+
+transform_kernel_allowlist: TypeAlias = reduction_kernels | transformation_kernels

pandas-stubs/core/groupby/generic.pyi

@@ -19,6 +19,7 @@ from typing import (
 from matplotlib.axes import Axes as PlotAxes
 import numpy as np
 from pandas.core.frame import DataFrame
+from pandas.core.groupby.base import transform_kernel_allowlist
 from pandas.core.groupby.groupby import (
     GroupBy,
     GroupByPlot,
@@ -41,7 +42,6 @@ from pandas._typing import (
     ByT,
     CorrelationMethod,
     Dtype,
-    GroupByFuncStrs,
     IndexLabel,
     Level,
     ListLike,
@@ -109,7 +109,9 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
         **kwargs: Any,
     ) -> UnknownSeries: ...
     @overload
-    def transform(self, func: GroupByFuncStrs, *args, **kwargs) -> UnknownSeries: ...
+    def transform(
+        self, func: transform_kernel_allowlist, *args, **kwargs
+    ) -> UnknownSeries: ...
     def filter(
         self, func: Callable | str, dropna: bool = ..., *args, **kwargs
     ) -> Series: ...
@@ -253,7 +255,9 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
         **kwargs: Any,
     ) -> DataFrame: ...
     @overload
-    def transform(self, func: GroupByFuncStrs, *args, **kwargs) -> DataFrame: ...
+    def transform(
+        self, func: transform_kernel_allowlist, *args, **kwargs
+    ) -> DataFrame: ...
     def filter(
         self, func: Callable, dropna: bool = ..., *args, **kwargs
     ) -> DataFrame: ...