pandas-dev / pandas-stubs

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

`GroupBy.pipe()` missing types in arguments and correct return type #738

Closed Dr-Irv closed 9 months ago

Dr-Irv commented 1 year ago

In core/generic.pyi, we have for pipe():

def pipe(
        self, func: Callable[..., T] | tuple[Callable[..., T], str], *args, **kwargs
    ) -> T: ...

But in core/groupby/groupby.pyi, we have:

    def pipe(self, func: Callable, *args, **kwargs): ...

So we should investigate the proper definition for the groupby() version to determine the return type.

Sample code to illustrate issue:

from typing import Union, cast
from typing_extensions import reveal_type
import pandas as pd
from pandas.core.groupby.generic import DataFrameGroupBy

df = pd.DataFrame({"x": ["a", "b", "b", "a", "a", "c"], "y": [1, 2, 3, 4, 5, 6]})
df.index = pd.Index(df["y"] * 10)

def foo(g: Union[pd.DataFrame, DataFrameGroupBy]) -> pd.Series:
    return cast(pd.Series, g.sum())

reveal_type(df.pipe(foo))

reveal_type(df.groupby("x").pipe(foo))

Runtime reveals:

Runtime type is 'Series'
Runtime type is 'DataFrame'

pyright output:

gpipe.py:14:13 - information: Type of "df.pipe(foo)" is "Series[Unknown]"
gpipe.py:16:13 - information: Type of "df.groupby("x").pipe(foo)" is "Unknown"
kshitiz305 commented 1 year ago

Hi @Dr-Irv I would like to contribute to this.

Dr-Irv commented 1 year ago

Hi @Dr-Irv I would like to contribute to this.

I've assigned it to you.