GH456 First attempt GroupBy.transform improved typing

Author: loicdiridollou

pandas-stubs/_typing.pyi

@@ -925,6 +925,44 @@ 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/generic.pyi

@@ -7,6 +7,7 @@ from collections.abc import (
 )
 from typing import (
     Any,
+    Concatenate,
     Generic,
     Literal,
     NamedTuple,
@@ -31,15 +32,18 @@ from typing_extensions import (
 from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._typing import (
     S1,
+    S2,
     AggFuncTypeBase,
     AggFuncTypeFrame,
     ByT,
     CorrelationMethod,
     Dtype,
+    GroupByFuncStrs,
     IndexLabel,
     Level,
     ListLike,
     NsmallestNlargestKeep,
+    P,
     Scalar,
     TakeIndexer,
     WindowingEngine,
@@ -72,14 +76,24 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
         **kwargs,
     ) -> Series: ...
     agg = aggregate
+    @overload
     def transform(
         self,
-        func: Callable | str,
-        *args,
+        func: Callable[Concatenate[Series[S1], P], Series[S2]],
+        *args: Any,
         engine: WindowingEngine = ...,
         engine_kwargs: WindowingEngineKwargs = ...,
-        **kwargs,
+        **kwargs: Any,
+    ) -> Series[S2]: ...
+    @overload
+    def transform(
+        self,
+        func: Callable,
+        *args: Any,
+        **kwargs: Any,
     ) -> Series: ...
+    @overload
+    def transform(self, func: GroupByFuncStrs, *args, **kwargs) -> Series: ...
     def filter(
         self, func: Callable | str, dropna: bool = ..., *args, **kwargs
     ) -> Series: ...
@@ -206,14 +220,24 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
         **kwargs,
     ) -> DataFrame: ...
     agg = aggregate
+    @overload
     def transform(
         self,
-        func: Callable | str,
-        *args,
+        func: Callable[Concatenate[DataFrame, P], DataFrame],
+        *args: Any,
         engine: WindowingEngine = ...,
         engine_kwargs: WindowingEngineKwargs = ...,
-        **kwargs,
+        **kwargs: Any,
     ) -> DataFrame: ...
+    @overload
+    def transform(
+        self,
+        func: Callable,
+        *args: Any,
+        **kwargs: Any,
+    ) -> DataFrame: ...
+    @overload
+    def transform(self, func: GroupByFuncStrs, *args, **kwargs) -> DataFrame: ...
     def filter(
         self, func: Callable, dropna: bool = ..., *args, **kwargs
     ) -> DataFrame: ...

tests/test_series.py

@@ -1084,6 +1084,24 @@ def test_types_groupby_agg() -> None:
         )
 
 
+def test_types_groupby_transform() -> None:
+    s: pd.Series[int] = pd.Series([4, 2, 1, 8], index=["a", "b", "a", "b"])
+
+    def transform_func(
+        x: pd.Series[int], pos_arg: bool, kw_arg: str
+    ) -> pd.Series[float]:
+        return x / (2.0 if pos_arg else 1.0)
+
+    check(
+        assert_type(
+            s.groupby(lambda x: x).transform(transform_func, True, kw_arg="foo"),
+            "pd.Series[float]",
+        ),
+        pd.Series,
+        float,
+    )
+
+
 def test_types_groupby_aggregate() -> None:
     s = pd.Series([4, 2, 1, 8], index=["a", "b", "a", "b"])
     check(assert_type(s.groupby(level=0).aggregate("sum"), pd.Series), pd.Series)