Repository: pandas-dev/pandas-stubs (language: python)

Issue:
BUG: df.groupby(col, as_index=False).value_counts() returns a DataFrame but is annotated as Series
**Describe the bug**

Using 
```
df = pd.DataFrame({"a": [1, 2, 3], "b": 1})

x: pd.DataFrame = df.groupby("a", as_index=False).value_counts()
```
raises an error, because mypy thinks the return is a Series

```
 Incompatible types in assignment (expression has type "Series[int]", variable has type "DataFrame")  [assignment]
``` 

**To Reproduce**
1. Provide a minimal runnable `pandas` example that is not properly checked by the stubs.
2. Indicate which type checker you are using (`mypy` or  `pyright`).
3. Show the error message received from that type checker while checking your example.

```
df = pd.DataFrame({"a": [1, 2, 3], "b": 1})

x: pd.DataFrame = df.groupby("a", as_index=False).value_counts()
```

In theory this should be ok, but not sure if this ``as_index=False`` can be typed?

**Please complete the following information:**
 - OS:  ubuntu
 - OS Version 22.04
 - python version 3.10
 - version of type checker mypy 0.971
 - version of installed `pandas-stubs` newest 20220815


**Additional context**
Add any other context about the problem here.



Discussion / hints:
Similar issue with `df.groupby('a",as_index=False).size()`, which returns a `DataFrame`, but we type it as `Series`
Yep, this happens probably for every function where a Series is possible
Two options here, one easy and one hard:
1. Easy way - just have `value_counts()` typed as `FrameOrSeriesUnion` as the result.
2. Hard way  (and may not be possible) - make `DataFrameGroupBy` and `SeriesGroupBy` Generic with a typing parameter that indices whether `Series` or `DataFrame` should be returned, based on the value of `is_index` in `DataFrame.groupby()` and `Series.groupby()`.  Requires overloads there, and careful typing.

Regarding (2), I recently did something like this to handle the case of `Series.str.split()` where the parameter `expand` changes the nature of the result.

If doable, I would obviously prefer 2 :)
I think `as_index` is kind of special and changes the return type of many functions, so there is also a third option of breaking `SeriesGroupBy` and `DataFrameGroupBy` into 2 different classes depending on `as_index` and overloading the various functions which return them.
> there is also a third option of breaking `SeriesGroupBy` and `DataFrameGroupBy` into 2 different classes depending on `as_index` and overloading the various functions which return them.

This is similar to my (2), but instead of using 4 total classes, we'd have 2 classes, each generic, and then the "parameter" to the generic class would indicate the return type to use based on `as_index`.  Need to explore both approaches.
I'm facing a case that type of `df.groupby("col1")["col2"].value_counts()` (with default of `as_index=True`) is `DataFrame`.
Tracing to the type definitions this call is using `SeriesGroupBy.value_counts()`.
I also checked in REPL that the return is definitely `Series` (with `MultiIndex`).

It seems to me that the definition look up is assuming `as_index=False`. Or is it something I do that confused it?
> I'm facing a case that type of `df.groupby("col1")["col2"].value_counts()` (with default of `as_index=True`) is `DataFrame`. Tracing to the type definitions this call is using `SeriesGroupBy.value_counts()`. I also checked in REPL that the return is definitely `Series` (with `MultiIndex`).
> 
> It seems to me that the definition look up is assuming `as_index=False`. Or is it something I do that confused it?

Can you provide a complete example?

![Screenshot from 2022-11-23 13-31-44](https://user-images.githubusercontent.com/544400/203476010-41e159b9-5e6e-4d41-b3d8-1fd12a4ddcfd.png)

```py
#!/usr/bin/env python3
import pandas as pd

df = pd.DataFrame(
    {
        "Animal": ["Falcon", "Falcon", "Parrot", "Parrot"],
        "Max Speed": [380, 370, 24, 26],
    }
)
c: pd.Series = df.groupby("Animal")["Max Speed"].value_counts()
print(type(c))  # <class 'pandas.core.series.Series'>
print(c.index)
print(c)
```

I concur with @amotzop's suggested solution.
> 'm facing a case that type of `df.groupby("col1")["col2"].value_counts()` (with default of `as_index=True`) is `DataFrame`.
> Tracing to the type definitions this call is using `SeriesGroupBy.value_counts()`.
> I also checked in REPL that the return is definitely `Series` (with `MultiIndex`).

This is a separate issue from the `as_index=False` issue, so created  #442
Thanks, this issues was the closest I got from Googling so I chime in here. :stuck_out_tongue_closed_eyes: 
@Dr-Irv looks like no work has been pushed on this and I'm quite keen on taking it.

I have encountered this issue from both the DataFrameGroupBy side as described above and the SeriesGroupBy as in 553.

I checked what you did on `Series.str.split` as you mention above, seems like the most elegant way to tackle this. 

I will try to send a PR in the coming week.
> @Dr-Irv looks like no work has been pushed on this and I'm quite keen on taking it.
> 
> I have encountered this issue from both the DataFrameGroupBy side as described above and the SeriesGroupBy as in 553.
> 
> I checked what you did on Series.str.split as you mention above, seems like the most elegant way to tackle this.
> 
> I will try to send a PR in the coming week.

Great!

Just note that since I suggested the solution, the `groupby` definitions in `series.pyi` and `frame.pyi` have been split into 8 overloads each, dependent on the type of the `by` argument.  So you'd end up with 16 overloads now, corresponding to the various combinations of the `by` and `as_index` arguments, and both `DataFrameGroupBy` and `SeriesGroupBy` would each add a new argument to the `Generic` declaration that would correspond to the return type needed for `value_counts()` (and `size()`, and maybe other methods as well).

It becomes a bit similar to what is in [pandas-stubs/core/indexes/accessors.pyi](https://github.com/pandas-dev/pandas-stubs/blob/main/pandas-stubs/core/indexes/accessors.pyi), which is used by `Series.dt`, `Timestamp`, `DatetimeIndex` and other places to have a consistent way of defining various time-related methods.



Fix patch:
diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
index ec12ba00..62192478 100644
--- a/pandas-stubs/core/frame.pyi
+++ b/pandas-stubs/core/frame.pyi
@@ -1054,29 +1054,53 @@ class DataFrame(NDFrame, OpsMixin):
         errors: IgnoreRaise = ...,
     ) -> None: ...
     @overload
-    def groupby(
+    def groupby(  # type: ignore[overload-overlap] # pyright: ignore reportOverlappingOverload
         self,
         by: Scalar,
         axis: AxisIndex | NoDefault = ...,
         level: IndexLabel | None = ...,
-        as_index: _bool = ...,
+        as_index: Literal[True] = True,
         sort: _bool = ...,
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[Scalar]: ...
+    ) -> DataFrameGroupBy[Scalar, Literal[True]]: ...
     @overload
     def groupby(
+        self,
+        by: Scalar,
+        axis: AxisIndex | NoDefault = ...,
+        level: IndexLabel | None = ...,
+        as_index: Literal[False] = ...,
+        sort: _bool = ...,
+        group_keys: _bool = ...,
+        observed: _bool | NoDefault = ...,
+        dropna: _bool = ...,
+    ) -> DataFrameGroupBy[Scalar, Literal[False]]: ...
+    @overload
+    def groupby(  # type: ignore[overload-overlap] # pyright: ignore reportOverlappingOverload
         self,
         by: DatetimeIndex,
         axis: AxisIndex | NoDefault = ...,
         level: IndexLabel | None = ...,
-        as_index: _bool = ...,
+        as_index: Literal[True] = True,
+        sort: _bool = ...,
+        group_keys: _bool = ...,
+        observed: _bool | NoDefault = ...,
+        dropna: _bool = ...,
+    ) -> DataFrameGroupBy[Timestamp, Literal[True]]: ...
+    @overload
+    def groupby(  # type: ignore[overload-overlap]
+        self,
+        by: DatetimeIndex,
+        axis: AxisIndex | NoDefault = ...,
+        level: IndexLabel | None = ...,
+        as_index: Literal[False] = ...,
         sort: _bool = ...,
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[Timestamp]: ...
+    ) -> DataFrameGroupBy[Timestamp, Literal[False]]: ...
     @overload
     def groupby(
         self,
@@ -1088,7 +1112,7 @@ class DataFrame(NDFrame, OpsMixin):
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[Timedelta]: ...
+    ) -> DataFrameGroupBy[Timedelta, bool]: ...
     @overload
     def groupby(
         self,
@@ -1100,7 +1124,7 @@ class DataFrame(NDFrame, OpsMixin):
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[Period]: ...
+    ) -> DataFrameGroupBy[Period, bool]: ...
     @overload
     def groupby(
         self,
@@ -1112,7 +1136,7 @@ class DataFrame(NDFrame, OpsMixin):
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[IntervalT]: ...
+    ) -> DataFrameGroupBy[IntervalT, bool]: ...
     @overload
     def groupby(
         self,
@@ -1124,7 +1148,7 @@ class DataFrame(NDFrame, OpsMixin):
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[tuple]: ...
+    ) -> DataFrameGroupBy[tuple, bool]: ...
     @overload
     def groupby(
         self,
@@ -1136,7 +1160,7 @@ class DataFrame(NDFrame, OpsMixin):
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[SeriesByT]: ...
+    ) -> DataFrameGroupBy[SeriesByT, bool]: ...
     @overload
     def groupby(
         self,
@@ -1148,7 +1172,7 @@ class DataFrame(NDFrame, OpsMixin):
         group_keys: _bool = ...,
         observed: _bool | NoDefault = ...,
         dropna: _bool = ...,
-    ) -> DataFrameGroupBy[Any]: ...
+    ) -> DataFrameGroupBy[Any, bool]: ...
     def pivot(
         self,
         *,
diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi
index 3d43e75e..8f2b5588 100644
--- a/pandas-stubs/core/groupby/generic.pyi
+++ b/pandas-stubs/core/groupby/generic.pyi
@@ -11,6 +11,7 @@ from typing import (
     Generic,
     Literal,
     NamedTuple,
+    TypeVar,
     final,
     overload,
 )
@@ -29,6 +30,7 @@ from typing_extensions import (
 )
 
 from pandas._libs.lib import NoDefault
+from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._typing import (
     S1,
     AggFuncTypeBase,
@@ -182,7 +184,9 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
         self,
     ) -> Iterator[tuple[ByT, Series[S1]]]: ...
 
-class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT]):
+_TT = TypeVar("_TT", bound=Literal[True, False])
+
+class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
     # error: Overload 3 for "apply" will never be used because its parameters overlap overload 1
     @overload  # type: ignore[override]
     def apply(  # type: ignore[overload-overlap]
@@ -236,7 +240,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT]):
     @overload
     def __getitem__(  # pyright: ignore[reportIncompatibleMethodOverride, reportOverlappingOverload]
         self, key: Iterable[Hashable] | slice
-    ) -> DataFrameGroupBy[ByT]: ...
+    ) -> DataFrameGroupBy[ByT, bool]: ...
     def nunique(self, dropna: bool = ...) -> DataFrame: ...
     def idxmax(
         self,
@@ -388,3 +392,11 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT]):
     def __iter__(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
     ) -> Iterator[tuple[ByT, DataFrame]]: ...
+    @overload
+    def size(self: DataFrameGroupBy[ByT, Literal[True]]) -> Series[int]: ...
+    @overload
+    def size(self: DataFrameGroupBy[ByT, Literal[False]]) -> DataFrame: ...
+    @overload
+    def size(self: DataFrameGroupBy[Timestamp, Literal[True]]) -> Series[int]: ...
+    @overload
+    def size(self: DataFrameGroupBy[Timestamp, Literal[False]]) -> DataFrame: ...
diff --git a/pandas-stubs/core/groupby/groupby.pyi b/pandas-stubs/core/groupby/groupby.pyi
index a009e50c..66e6f825 100644
--- a/pandas-stubs/core/groupby/groupby.pyi
+++ b/pandas-stubs/core/groupby/groupby.pyi
@@ -232,11 +232,7 @@ class GroupBy(BaseGroupBy[NDFrameT]):
     def sem(
         self: GroupBy[DataFrame], ddof: int = ..., numeric_only: bool = ...
     ) -> DataFrame: ...
-    @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 | Series[int]: ...
     @final
     def sum(
         self,


Test patch:
diff --git a/tests/test_frame.py b/tests/test_frame.py
index 64198952..c75ecd12 100644
--- a/tests/test_frame.py
+++ b/tests/test_frame.py
@@ -1025,6 +1025,40 @@ def test_types_pivot_table() -> None:
     )
 
 
+def test_types_groupby_as_index() -> None:
+    df = pd.DataFrame({"a": [1, 2, 3]})
+    check(
+        assert_type(
+            df.groupby("a", as_index=False).size(),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            df.groupby("a", as_index=True).size(),
+            "pd.Series[int]",
+        ),
+        pd.Series,
+    )
+
+
+def test_types_groupby_size() -> None:
+    """Test for GH886."""
+    data = [
+        {"date": "2023-12-01", "val": 12},
+        {"date": "2023-12-02", "val": 2},
+        {"date": "2023-12-03", "val": 1},
+        {"date": "2023-12-03", "val": 10},
+    ]
+
+    df = pd.DataFrame(data)
+    groupby = df.groupby("date")
+    size = groupby.size()
+    frame = size.to_frame()
+    check(assert_type(frame.reset_index(), pd.DataFrame), pd.DataFrame)
+
+
 def test_types_groupby() -> None:
     df = pd.DataFrame(data={"col1": [1, 1, 2], "col2": [3, 4, 5], "col3": [0, 1, 0]})
     df.index.name = "ind"