Update to the fix

Author: loicdiridollou

pandas-stubs/core/frame.pyi

@@ -12,7 +12,10 @@ from re import Pattern
 from typing import (
     Any,
     ClassVar,
+    Generic,
     Literal,
+    TypeVar,
+    Union,
     overload,
 )
 
@@ -77,6 +80,7 @@ from pandas._typing import (
     Axis,
     AxisColumn,
     AxisIndex,
+    ByT,
     CalculationMethod,
     ColspaceArgType,
     CompressionOptions,
@@ -232,6 +236,14 @@ class _LocIndexerFrame(_LocIndexer):
         value: Scalar | NAType | NaTType | ArrayLike | Series | list | None,
     ) -> None: ...
 
+TT = TypeVar("TT", bound=Union[Literal[True], Literal[False]])
+
+class DataFrameGroupByGen(DataFrameGroupBy[ByT], Generic[ByT, TT]):
+    pass
+
+class SeriesGroupByGen(SeriesGroupBy, Generic[TT, ByT]):
+    pass
+
 class DataFrame(NDFrame, OpsMixin):
     __hash__: ClassVar[None]  # type: ignore[assignment]
 
@@ -1055,29 +1067,29 @@ class DataFrame(NDFrame, OpsMixin):
         errors: IgnoreRaise = ...,
     ) -> None: ...
     @overload
-    def groupby(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+    def groupby(  # type: ignore[overload-overlap] # pyright: ignore reportOverlappingOverload
         self,
         by: Scalar,
         axis: AxisIndex | NoDefault = ...,
         level: IndexLabel | None = ...,
-        as_index: Literal[False] = ...,
+        as_index: Literal[True] = True,
         sort: _bool = ...,
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[Scalar]: ...
+    ) -> DataFrameGroupByGen[Scalar, Literal[True]]: ...
     @overload
     def groupby(
         self,
         by: Scalar,
         axis: AxisIndex | NoDefault = ...,
         level: IndexLabel | None = ...,
-        as_index: Literal[True] = True,
+        as_index: Literal[False] = ...,
         sort: _bool = ...,
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> SeriesGroupBy: ...
+    ) -> DataFrameGroupByGen[Scalar, Literal[False]]: ...
     @overload
     def groupby(
         self,

pandas-stubs/core/groupby/groupby.pyi

@@ -18,7 +18,10 @@ from typing import (
 
 import numpy as np
 from pandas.core.base import SelectionMixin
-from pandas.core.frame import DataFrame
+from pandas.core.frame import (
+    DataFrame,
+    DataFrameGroupByGen,
+)
 from pandas.core.groupby import (
     generic,
     ops,
@@ -53,6 +56,7 @@ from pandas._typing import (
     AnyArrayLike,
     Axis,
     AxisInt,
+    ByT,
     CalculationMethod,
     Dtype,
     Frequency,
@@ -235,8 +239,10 @@ class GroupBy(BaseGroupBy[NDFrameT]):
     @final
     @overload
     def size(self: GroupBy[Series]) -> Series[int]: ...
-    @overload  # return type depends on `as_index` for dataframe groupby
-    def size(self: GroupBy[DataFrame]) -> DataFrame: ...
+    @overload
+    def size(self: DataFrameGroupByGen[ByT, Literal[True]]) -> Series[int]: ...  # type: ignore[misc]
+    @overload
+    def size(self: DataFrameGroupByGen[ByT, Literal[False]]) -> DataFrame: ...  # type: ignore[misc]
     @final
     def sum(
         self,

tests/test_frame.py

@@ -1066,7 +1066,7 @@ def test_types_groupby() -> None:
 
     df1: pd.DataFrame = df.groupby(by="col1").agg("sum")
     df2: pd.DataFrame = df.groupby(level="ind").aggregate("sum")
-    df3: pd.Series = df.groupby(by="col1", sort=False, as_index=True).transform(
+    df3: pd.DataFrame = df.groupby(by="col1", sort=False, as_index=True).transform(
         lambda x: x.max()
     )
     df4: pd.DataFrame = df.groupby(by=["col1", "col2"]).count()