BUG: df.groupby(col, as_index=False).value_counts() returns a DataFrame but is annotated as Series

[phofl]

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.

[Dr-Irv]

Similar issue with df.groupby('a",as_index=False).size(), which returns a DataFrame, but we type it as Series

[phofl]

Yep, this happens probably for every function where a Series is possible

[Dr-Irv]

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.

[phofl]

If doable, I would obviously prefer 2 :)

[amotzop]

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.

[Dr-Irv]

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.

[leesei]

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?

[Dr-Irv]

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?

[leesei]

#!/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.

[Dr-Irv]

'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

[leesei]

Thanks, this issues was the closest I got from Googling so I chime in here. 😝