Skip to content

GH456 First attempt GroupBy.transform improved typing #1242

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

GH456 First attempt GroupBy.transform improved typing #1242

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
020f93d
GH456 First attempt GroupBy.transform improved typing
loicdiridollou Jun 6, 2025
106a6f5
GH456 Attempt GroupBy.aggregate improved typing
loicdiridollou Jun 6, 2025
3bba101
GH456 Attempt GroupBy.aggregate improved typing
loicdiridollou Jun 6, 2025
053b7e7
GH456 PR Feedback
loicdiridollou Jun 7, 2025
4141a06
GH456 PR Feedback
loicdiridollou Jun 7, 2025
f9863d0
GH456 PR Feedback
loicdiridollou Jun 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 49 additions & 0 deletions pandas-stubs/core/groupby/base.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,56 @@
from collections.abc import Hashable
import dataclasses
from typing import (
Literal,
TypeAlias,
)

@dataclasses.dataclass(order=True, frozen=True)
class OutputKey:
label: Hashable
position: int

ReductionKernelType: TypeAlias = Literal[
"all",
"any",
"corrwith",
"count",
"first",
"idxmax",
"idxmin",
"last",
"max",
"mean",
"median",
"min",
"nunique",
"prod",
# as long as `quantile`'s signature accepts only
# a single quantile value, it's a reduction.
# GH#27526 might change that.
"quantile",
"sem",
"size",
"skew",
"std",
"sum",
"var",
]

TransformationKernelType: TypeAlias = Literal[
"bfill",
"cumcount",
"cummax",
"cummin",
"cumprod",
"cumsum",
"diff",
"ffill",
"fillna",
"ngroup",
"pct_change",
"rank",
"shift",
]

TransformReductionListType: TypeAlias = ReductionKernelType | TransformationKernelType
71 changes: 62 additions & 9 deletions pandas-stubs/core/groupby/generic.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ from collections.abc import (
)
from typing import (
Any,
Concatenate,
Generic,
Literal,
NamedTuple,
Expand All @@ -18,11 +19,15 @@ from typing import (
from matplotlib.axes import Axes as PlotAxes
import numpy as np
from pandas.core.frame import DataFrame
from pandas.core.groupby.base import TransformReductionListType
from pandas.core.groupby.groupby import (
GroupBy,
GroupByPlot,
)
from pandas.core.series import Series
from pandas.core.series import (
Series,
UnknownSeries,
)
from typing_extensions import (
Self,
TypeAlias,
Expand All @@ -31,6 +36,7 @@ from typing_extensions import (
from pandas._libs.tslibs.timestamps import Timestamp
from pandas._typing import (
S1,
S2,
AggFuncTypeBase,
AggFuncTypeFrame,
ByT,
Expand All @@ -40,6 +46,7 @@ from pandas._typing import (
Level,
ListLike,
NsmallestNlargestKeep,
P,
Scalar,
TakeIndexer,
WindowingEngine,
Expand All @@ -53,10 +60,30 @@ class NamedAgg(NamedTuple):
aggfunc: AggScalar

class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
@overload
def aggregate(
self,
func: Callable[Concatenate[Series[S1], P], S2],
/,
*args,
engine: WindowingEngine = ...,
engine_kwargs: WindowingEngineKwargs = ...,
**kwargs,
) -> Series[S2]: ...
@overload
def aggregate(
self,
func: Callable[[Series], S2],
*args,
engine: WindowingEngine = ...,
engine_kwargs: WindowingEngineKwargs = ...,
**kwargs,
) -> Series[S2]: ...
@overload
def aggregate(
self,
func: list[AggFuncTypeBase],
/,
*args,
engine: WindowingEngine = ...,
engine_kwargs: WindowingEngineKwargs = ...,
Expand All @@ -66,20 +93,34 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
def aggregate(
self,
func: AggFuncTypeBase | None = ...,
/,
*args,
Comment on lines 93 to 97
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Before this overload, you could add this overload:

    @overload
    def aggregate(
        self,
        func: Callable[[Series], S2],
        *args,
        engine: WindowingEngine = ...,
        engine_kwargs: WindowingEngineKwargs = ...,
        **kwargs,
    ) -> Series[S2]: ...

Then you know that if you start with a Series with a known type, then the return type would be inferred from the callable. And it works with a lambda function, e.g.:

    s = pd.Series([1, 2, 3, 4])
    q = s.groupby([1, 1, 2, 2]).agg(lambda x: x.astype(float).min())

In this case, q would have type Series[float], which is what you want.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

image Even with the new overload mypy still complains (pyright does not). I think it does not recognize the lamda as returning S2 (float).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that's because the type of new_func isn't clear.

But I think it would work if you did check(assert_type(s.groupby([1,1,2,2]).agg(lambda x: x.astype(float).min()), "pd.Series[int]"), pd.Series, int)

Because then it can know that x is a Series[int] and that the lambda becomes Series[int]

Can you try that?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried that for the last push, see

pandas-stubs/tests/test_series.py

Line 1167 in f9863d0

check(assert_type(s.groupby([1,1,2,2]).agg(lambda x: x.astype(float).min()), "pd.Series[float]"), pd.Series, int)

It fails in all CI:

===========================================
Beginning: 'Run mypy on 'tests' (using the local stubs) and on the local stubs'
===========================================

tests/test_series.py:1167: error: Expression is of type "Series[Any]", not "Series[float]"  [assert-type]
Found 1 error in 1 file (checked 224 source files)

engine: WindowingEngine = ...,
engine_kwargs: WindowingEngineKwargs = ...,
**kwargs,
) -> Series: ...
) -> UnknownSeries: ...
agg = aggregate
@overload
def transform(
self,
func: Callable | str,
*args,
func: Callable[Concatenate[Series[S1], P], Series[S2]],
/,
*args: Any,
engine: WindowingEngine = ...,
engine_kwargs: WindowingEngineKwargs = ...,
**kwargs,
) -> Series: ...
**kwargs: Any,
) -> Series[S2]: ...
@overload
def transform(
self,
func: Callable,
*args: Any,
**kwargs: Any,
) -> UnknownSeries: ...
@overload
def transform(
self, func: TransformReductionListType, *args, **kwargs
) -> UnknownSeries: ...
def filter(
self, func: Callable | str, dropna: bool = ..., *args, **kwargs
) -> Series: ...
Expand Down Expand Up @@ -206,13 +247,25 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
**kwargs,
) -> DataFrame: ...
agg = aggregate
@overload
def transform(
self,
func: Callable | str,
*args,
func: Callable[Concatenate[DataFrame, P], DataFrame],
*args: Any,
engine: WindowingEngine = ...,
engine_kwargs: WindowingEngineKwargs = ...,
**kwargs,
**kwargs: Any,
) -> DataFrame: ...
@overload
def transform(
self,
func: Callable,
*args: Any,
**kwargs: Any,
) -> DataFrame: ...
@overload
def transform(
self, func: TransformReductionListType, *args, **kwargs
) -> DataFrame: ...
def filter(
self, func: Callable, dropna: bool = ..., *args, **kwargs
Expand Down
88 changes: 86 additions & 2 deletions tests/test_series.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1078,25 +1078,109 @@ def test_types_groupby_agg() -> None:
r"The provided callable <built-in function (min|sum)> is currently using",
upper="2.2.99",
):
check(assert_type(s.groupby(level=0).agg(sum), pd.Series), pd.Series)

def sum_sr(s: pd.Series[int]) -> int:
# type of `sum` not well inferred by mypy
return s.sum()

check(
assert_type(s.groupby(level=0).agg(sum_sr), "pd.Series[int]"),
pd.Series,
np.integer,
)
check(
assert_type(s.groupby(level=0).agg([min, sum]), pd.DataFrame), pd.DataFrame
)


def test_types_groupby_transform() -> None:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you should add tests for two of the string transform arguments (e.g., "mean", "first")

s: pd.Series[int] = pd.Series([4, 2, 1, 8], index=["a", "b", "a", "b"])

def transform_func(
x: pd.Series[int], pos_arg: bool, kw_arg: str
) -> pd.Series[float]:
return x / (2.0 if pos_arg else 1.0)

check(
assert_type(
s.groupby(lambda x: x).transform(transform_func, True, kw_arg="foo"),
"pd.Series[float]",
),
pd.Series,
float,
)
check(
assert_type(
s.groupby(lambda x: x).transform(
transform_func, True, engine="cython", kw_arg="foo"
),
"pd.Series[float]",
),
pd.Series,
float,
)
check(
assert_type(
s.groupby(lambda x: x).transform("mean"),
"pd.Series",
),
pd.Series,
)
check(
assert_type(
s.groupby(lambda x: x).transform("first"),
"pd.Series",
),
pd.Series,
)


def test_types_groupby_aggregate() -> None:
s = pd.Series([4, 2, 1, 8], index=["a", "b", "a", "b"])
check(assert_type(s.groupby(level=0).aggregate("sum"), pd.Series), pd.Series)
check(
assert_type(s.groupby(level=0).aggregate(["min", "sum"]), pd.DataFrame),
pd.DataFrame,
)

def func(s: pd.Series[int]) -> float:
return s.astype(float).min()

s = pd.Series([1, 2, 3, 4])
check(
assert_type(s.groupby([1, 1, 2, 2]).agg(func), "pd.Series[float]"),
pd.Series,
np.floating,
)
check(
assert_type(s.groupby(level=0).aggregate(func), "pd.Series[float]"),
pd.Series,
np.floating,
)
check(
assert_type(
s.groupby(level=0).aggregate(func, engine="cython"), "pd.Series[float]"
),
pd.Series,
np.floating,
)
check(assert_type(s.groupby([1,1,2,2]).agg(lambda x: x.astype(float).min()), "pd.Series[float]"), pd.Series, int)

with pytest_warns_bounded(
FutureWarning,
r"The provided callable <built-in function (min|sum)> is currently using",
upper="2.2.99",
):
check(assert_type(s.groupby(level=0).aggregate(sum), pd.Series), pd.Series)

def sum_sr(s: pd.Series[int]) -> int:
# type of `sum` not well inferred by mypy
return s.sum()

check(
assert_type(s.groupby(level=0).aggregate(sum_sr), "pd.Series[int]"),
pd.Series,
np.integer,
)
check(
assert_type(s.groupby(level=0).aggregate([min, sum]), pd.DataFrame),
pd.DataFrame,
Expand Down
Loading