GroupBy[Series].count() should have return type Series[int]

pandas-dev / pandas-stubs

Public type stubs for pandas

BSD 3-Clause "New" or "Revised" License

229 stars 122 forks source link

GroupBy[Series].count() should have return type Series[int] #965

Closed chrisyeh96 closed 1 month ago

chrisyeh96 commented 1 month ago

Describe the bug

A clear and concise description of what the bug is.

To Reproduce

s = pd.Series(["a", "b", "a"], index=[1, 2, 2])
s.groupby(level=0).count()

Type checker: pyright (in VS Code)

The expected type is Series[int], but instead pyright shows it as Series[str]

Please complete the following information:

OS: Ubuntu
OS Version: 22.04.4 LTS
python version: 3.10.12
version of type checker: pyright 1.1.369
version of installed pandas-stubs: 2.2.2.240603

Dr-Irv commented 1 month ago

Thanks for the report. Need to remove def count() from core/groupby/groupby.pyi and then restate it in core/groupby/generic.pyi in both the DataFrameGroupBy and SeriesGroupBy classes.

PR with tests welcome.