pandas-dev / pandas-stubs

Public type stubs for pandas
BSD 3-Clause "New" or "Revised" License
236 stars 129 forks source link

Changes to `groupby.size()` don't work for all possible groupby #1045

Closed Dr-Irv closed 3 days ago

Dr-Irv commented 3 days ago

Describe the bug

import pandas as pd
ddf = pd.DataFrame({"a": [1, 1, 2], "b": [2, 3, 2]})
reveal_type(ddf.groupby(["a", "b"]).size())

pyright reports

 - information: Type of "ddf.groupby(["a", "b"]).size()" is "Unknown"
  c:\Code\pandas-stubs\play\issue1044.py:3:37 - error: Cannot access attribute "size" for class "DataFrameGroupBy[tuple[Unknown, ...], bool]"
    Could not bind method "size" because "DataFrameGroupBy[tuple[Unknown, ...], bool]" is not assignable to parameter "self"
      "DataFrameGroupBy[tuple[Unknown, ...], bool]" is not assignable to "DataFrameGroupBy[tuple[Unknown, ...], Literal[True]]"
        Type parameter "_TT@DataFrameGroupBy" is invariant, but "bool" is not the same as "Literal[True]"
    Could not bind method "size" because "DataFrameGroupBy[tuple[Unknown, ...], bool]" is not assignable to parameter "self"
      "DataFrameGroupBy[tuple[Unknown, ...], bool]" is not assignable to "DataFrameGroupBy[tuple[Unknown, ...], Literal[False]]"
        Type parameter "_TT@DataFrameGroupBy" is invariant, but "bool" is not the same as "Literal[False]"
    Could not bind method "size" because "DataFrameGroupBy[tuple[Unknown, ...], bool]" is not assignable to parameter "self"
      "DataFrameGroupBy[tuple[Unknown, ...], bool]" is not assignable to "DataFrameGroupBy[Timestamp, Literal[True]]" (reportAttributeAccessIssue)

mypy reports

note: Revealed type is "Any"

The fix here is that any of the methods in frame.pyi for def groupby() that return DataFrameGroupBy[SOMETYPE, bool] need to be split into 2 where there is one overload that returns DataFrameGroupBy[SOMETYPE, True] and another that returns DataFrameGroupBy[SOMETYPE, False] like is done with DataFrameGroupBy[Scalar, True] and DataFrameGroupBy[Scalar, False] .

Please complete the following information:

Additional context

This was missed in PR #1014 by @loicdiridollou . Hopefully he can fix soon, since the prerelease of VS Code pylance has this bug.

loicdiridollou commented 3 days ago

Sorry for that one, PR should be up.