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

Merged
merged 8 commits into from
Jun 13, 2025
Merged

GH456 First attempt GroupBy.transform improved typing #1242

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
8 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
e26b4c1
GH456 PR Feedback
loicdiridollou Jun 10, 2025
96abf3b
GH456 PR Feedback
loicdiridollou Jun 11, 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)

Copy link
Member Author

Choose a reason for hiding this comment

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

When I look with how mypy reads the type of the lambda, it has no idea about the type of x:

tests/test_series.py:1168: note: Revealed type is "def (x: Any) -> Any"

so that may explain why it fails on lambda expressions whatsoever.

Copy link
Collaborator

Choose a reason for hiding this comment

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

OK - so we can leave the lambda test in, but just have it assert_type() against Series instead of Series[float]

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
90 changes: 88 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,111 @@ 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,
)

# test below passes with mypy but pyright correctly sees it as pd.Series[float]
Copy link
Collaborator

Choose a reason for hiding this comment

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

Just have to change the comment to say "fails with mypy"

# check(assert_type(s.groupby([1,1,2,2]).agg(lambda x: x.astype(float).min()), pd.Series), pd.Series, float)
Copy link
Collaborator

Choose a reason for hiding this comment

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

Keep the commented test in there so it is still there and executes, since it works for both type checkers, but comment out the one that is "better" that has pyright infer it as Series[float].

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 am forced to comment it out because pyright sees it as pd.Series[float] but mypy sees it as pd.Series so in both versions the CI will fail either for mypy or pyright step. What would you recommend doing?

Copy link
Collaborator

Choose a reason for hiding this comment

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

I guess you have to keep it commented out. Do you have a test like this that passes both checkers:

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

So you can have the "preferred" version in there commented out, but I think the above test would pass both type checkers.

Copy link
Member Author

Choose a reason for hiding this comment

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

Actually that also fails to pass with mypy (pyright is fine with it).

Copy link
Collaborator

Choose a reason for hiding this comment

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

I tried a bunch of ideas and couldn't get it to work. It's probably a mypy bug, but I couldn't come up with a simple example that illustrates the problem.


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