GH456 Attempt GroupBy.aggregate improved typing

Author: loicdiridollou

pandas-stubs/core/groupby/generic.pyi

@@ -23,7 +23,10 @@ from pandas.core.groupby.groupby import (
     GroupBy,
     GroupByPlot,
 )
-from pandas.core.series import Series
+from pandas.core.series import (
+    Series,
+    UnknownSeries,
+)
 from typing_extensions import (
     Self,
     TypeAlias,
@@ -57,10 +60,21 @@ class NamedAgg(NamedTuple):
     aggfunc: AggScalar
 
 class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
+    @overload
+    def aggregate(
+        self,
+        func: Callable[Concatenate[Series[S1], P], S2],
+        /,
+        *args,
+        engine: WindowingEngine = ...,
+        engine_kwargs: WindowingEngineKwargs = ...,
+        **kwargs,
+    ) -> Series[S2]: ...
     @overload
     def aggregate(
         self,
         func: list[AggFuncTypeBase],
+        /,
         *args,
         engine: WindowingEngine = ...,
         engine_kwargs: WindowingEngineKwargs = ...,
@@ -70,16 +84,18 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
     def aggregate(
         self,
         func: AggFuncTypeBase | None = ...,
+        /,
         *args,
         engine: WindowingEngine = ...,
         engine_kwargs: WindowingEngineKwargs = ...,
         **kwargs,
-    ) -> Series: ...
+    ) -> UnknownSeries: ...
     agg = aggregate
     @overload
     def transform(
         self,
         func: Callable[Concatenate[Series[S1], P], Series[S2]],
+        /,
         *args: Any,
         engine: WindowingEngine = ...,
         engine_kwargs: WindowingEngineKwargs = ...,
@@ -91,9 +107,9 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
         func: Callable,
         *args: Any,
         **kwargs: Any,
-    ) -> Series: ...
+    ) -> UnknownSeries: ...
     @overload
-    def transform(self, func: GroupByFuncStrs, *args, **kwargs) -> Series: ...
+    def transform(self, func: GroupByFuncStrs, *args, **kwargs) -> UnknownSeries: ...
     def filter(
         self, func: Callable | str, dropna: bool = ..., *args, **kwargs
     ) -> Series: ...

tests/test_series.py

@@ -1078,7 +1078,16 @@ def test_types_groupby_agg() -> None:
         r"The provided callable <built-in function (min|sum)> is currently using",
         upper="2.2.99",
     ):
-        check(assert_type(s.groupby(level=0).agg(sum), pd.Series), pd.Series)
+
+        def sum_sr(s: pd.Series[int]) -> int:
+            # type of `sum` not well inferred by mypy
+            return sum(s)
+
+        check(
+            assert_type(s.groupby(level=0).agg(sum_sr), "pd.Series[int]"),
+            pd.Series,
+            np.integer,
+        )
         check(
             assert_type(s.groupby(level=0).agg([min, sum]), pd.DataFrame), pd.DataFrame
         )
@@ -1100,6 +1109,16 @@ def transform_func(
         pd.Series,
         float,
     )
+    check(
+        assert_type(
+            s.groupby(lambda x: x).transform(
+                transform_func, True, engine="cython", kw_arg="foo"
+            ),
+            "pd.Series[float]",
+        ),
+        pd.Series,
+        float,
+    )
 
 
 def test_types_groupby_aggregate() -> None:
@@ -1109,12 +1128,40 @@ def test_types_groupby_aggregate() -> None:
         assert_type(s.groupby(level=0).aggregate(["min", "sum"]), pd.DataFrame),
         pd.DataFrame,
     )
+
+    def func(s: pd.Series[int]) -> float:
+        return s.astype(float).min()
+
+    s = pd.Series([1, 2, 3, 4])
+    s.groupby([1, 1, 2, 2]).agg(lambda x: x.astype(float).min())
+    check(
+        assert_type(s.groupby(level=0).aggregate(func), "pd.Series[float]"),
+        pd.Series,
+        np.floating,
+    )
+    check(
+        assert_type(
+            s.groupby(level=0).aggregate(func, engine="cython"), "pd.Series[float]"
+        ),
+        pd.Series,
+        np.floating,
+    )
+
     with pytest_warns_bounded(
         FutureWarning,
         r"The provided callable <built-in function (min|sum)> is currently using",
         upper="2.2.99",
     ):
-        check(assert_type(s.groupby(level=0).aggregate(sum), pd.Series), pd.Series)
+
+        def sum_sr(s: pd.Series[int]) -> int:
+            # type of `sum` not well inferred by mypy
+            return sum(s)
+
+        check(
+            assert_type(s.groupby(level=0).aggregate(sum_sr), "pd.Series[int]"),
+            pd.Series,
+            np.integer,
+        )
         check(
             assert_type(s.groupby(level=0).aggregate([min, sum]), pd.DataFrame),
             pd.DataFrame,