Add overload for DataFrameGroupBy.groupby("size") return Series (#739)

* Add overload for DataFrameGroupBy.groupby("size") return Series

* Switch SeriesGroupBy.agg to assignment

* Move tests to test_types_groupby_agg

* Change error comment to different return types

Author: ljmc-github

pandas-stubs/core/groupby/generic.pyi

@@ -54,10 +54,7 @@ class SeriesGroupBy(GroupBy, Generic[S1, ByT]):
     def aggregate(self, func: list[AggFuncTypeBase], *args, **kwargs) -> DataFrame: ...
     @overload
     def aggregate(self, func: AggFuncTypeBase, *args, **kwargs) -> Series: ...
-    @overload
-    def agg(self, func: list[AggFuncTypeBase], *args, **kwargs) -> DataFrame: ...
-    @overload
-    def agg(self, func: AggFuncTypeBase, *args, **kwargs) -> Series: ...
+    agg = aggregate
     def transform(self, func: Callable | str, *args, **kwargs) -> Series: ...
     def filter(self, func, dropna: bool = ..., *args, **kwargs): ...
     def nunique(self, dropna: bool = ...) -> Series: ...
@@ -159,6 +156,10 @@ class DataFrameGroupBy(GroupBy, Generic[ByT]):
     def apply(  # pyright: ignore[reportOverlappingOverload]
         self, func: Callable[[Iterable], float], *args, **kwargs
     ) -> DataFrame: ...
+    # error: overload 1 overlaps overload 2 because of different return types
+    @overload
+    def aggregate(self, arg: Literal["size"]) -> Series: ...  # type: ignore[misc]  # pyright: ignore[reportOverlappingOverload]
+    @overload
     def aggregate(self, arg: AggFuncTypeFrame = ..., *args, **kwargs) -> DataFrame: ...
     agg = aggregate
     def transform(self, func: Callable | str, *args, **kwargs) -> DataFrame: ...

tests/test_frame.py

@@ -951,6 +951,9 @@ def wrapped_min(x: Any) -> Any:
 
     cols_mixed: list[str | int] = ["col1", 0]
     check(assert_type(df.groupby(by=cols_mixed).sum(), pd.DataFrame), pd.DataFrame)
+    # GH 736
+    check(assert_type(df.groupby(by="col1").aggregate("size"), pd.Series), pd.Series)
+    check(assert_type(df.groupby(by="col1").agg("size"), pd.Series), pd.Series)
 
 
 # This was added in 1.1.0 https://pandas.pydata.org/docs/whatsnew/v1.1.0.html