Unpacking group keys is refused in multi-columns groups

[mdespriee]

Describe the bug
Unpacking group keys is refused in multi-columns groups, group key is wrongly assumed to be a str

To Reproduce

import pandas as pd
df = pd.DataFrame({"a": [0, 1, 2], "b": [4, 5, 6], "c": [7, 8, 9]})
for (k1, k2), g in df.groupby(["a", "b"]):
    print(k1, k2)

Script output

0 4
1 5
2 6

mypy output

foo.py:5: error: Unpacking a string is disallowed
foo.py:6: error: Cannot determine type of "k1"
foo.py:6: error: Cannot determine type of "k2"
Found 3 errors in 1 file (checked 1 source file)

Please complete the following information:

• OS: Ubuntu 22.04
• python=3.9.13
• pandas=1.4.3
• pandas-stubs=1.4.3.220710
• mypy=0.961

[Dr-Irv]

Issue is in definition of __iter__() in pandas-stubs/core/groupby/groupby.pyi. Should use what is in pandas source which is Iterator[tuple[Hashable, NDFrameT]] as the return type.

PR welcome.

[Dr-Irv]

We can't really support this with static typing, because the iterator produces a Tuple[Hashable, NDFrameT] . The issue here is that the first component of the tuple could be another tuple (as in the case here) or could be a singular value.

>>> df = pd.DataFrame({"a": [0, 1, 2], "b": [4, 5, 6], "c": [7, 8, 9]})
>>> for k, g in df.groupby(["a", "b"]):
...     print(k, type(k))
...
(0, 4) <class 'tuple'>
(1, 5) <class 'tuple'>
(2, 6) <class 'tuple'>
>>> for k, g in df.groupby("a"):
...     print(k, type(k))
...
0 <class 'int'>
1 <class 'int'>
2 <class 'int'>

Since the groupby argument directly affects what the iterator produces, we can't track that.

The PR I created #148 does address partially this issue by improving the return type of __iter__(), but to do what you want here with respect to typing would need to be done as follows:

for k, g in df.groupby(["a", "b"]):
    k1, k2 = cast(Tuple[int, int], k)
    print(k1, k2)

[twoertwein]

Could have something like the following (I don't think it is worth it and it would also only work with tuple):

class DataFrame:
    @overload
    def groupby(keys: tuple[Hashable]) -> GroupByClass[tuple[Hashable]]: ...
    @overload
    def groupby(keys: tuple[Hashable, Hashable]) -> GroupByClass[tuple[Hashable, Hashable], DataFrame]: ...
    # ...
    @overload
    def groupby(keys: tuple[Hashable, ...]) -> GroupByClass[tuple[Hashable, ...], DataFrame]: ...

TuplesT = TypeVar("TuplesT", bound=tuple)
class GroupByClass(Generic[TuplesT, NDFrameT]):
    def __iter__(self) -> Iterator[tuple[TuplesT, NDFrameT]]: ...

[Dr-Irv]

I did something close to the above idea. Might not catch some esoteric cases, but does differentiate between grouping by a list versus grouping by a single value. See 383e1bf