towards groupby

Author: cmp0xff

pandas-stubs/core/groupby/generic.pyi

@@ -30,6 +30,7 @@ from typing_extensions import override
 
 from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._typing import (
+    C2,
     S2,
     S3,
     AggFuncTypeBase,
@@ -69,22 +70,13 @@ class SeriesGroupBy(GroupBy[Series[S2]], Generic[S2, ByT]):
     @overload
     def aggregate(
         self,
-        func: Callable[Concatenate[Series[S2], P], S3],
+        func: Callable[[Series], C2] | Callable[Concatenate[Series[S2], P], C2],
         /,
         *args: Any,
         engine: WindowingEngine = ...,
         engine_kwargs: WindowingEngineKwargs = ...,
         **kwargs: Any,
-    ) -> Series[S3]: ...
-    @overload
-    def aggregate(
-        self,
-        func: Callable[[Series], S3],
-        *args: Any,
-        engine: WindowingEngine = ...,
-        engine_kwargs: WindowingEngineKwargs = ...,
-        **kwargs: Any,
-    ) -> Series[S3]: ...
+    ) -> Series[C2]: ...
     @overload
     def aggregate(
         self,

tests/frame/test_groupby.py

@@ -7,6 +7,7 @@
     Iterator,
 )
 from typing import (
+    TYPE_CHECKING,
     Any,
     assert_type,
 )
@@ -26,6 +27,9 @@
     check,
 )
 
+if TYPE_CHECKING:
+    from pandas._typing import S1
+
 
 def test_types_groupby_as_index() -> None:
     """Test type of groupby.size method depending on `as_index`."""
@@ -296,7 +300,7 @@ def test_types_groupby_agg() -> None:
     agg_dict1 = {"col2": "min", "col3": "max", 0: "sum"}
     check(assert_type(df.groupby("col1").agg(agg_dict1), pd.DataFrame), pd.DataFrame)
 
-    def wrapped_min(x: pd.Series) -> Scalar:
+    def wrapped_min(x: pd.Series[S1]) -> S1:
         return x.min()  # type: ignore[no-any-return]
 
     # TODO: facebook/pyrefly#0