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 5 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
38 changes: 0 additions & 38 deletions pandas-stubs/_typing.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -925,44 +925,6 @@ GroupByObjectNonScalar: TypeAlias = (
| list[Grouper]
)
GroupByObject: TypeAlias = Scalar | Index | GroupByObjectNonScalar | Series
GroupByFuncStrs: TypeAlias = Literal[
# Reduction/aggregation functions
"all",
"any",
"corrwith",
"count",
"first",
"idxmax",
"idxmin",
"last",
"max",
"mean",
"median",
"min",
"nunique",
"prod",
"quantile",
"sem",
"size",
"skew",
"std",
"sum",
"var",
# Transformation functions
"bfill",
"cumcount",
"cummax",
"cummin",
"cumprod",
"cumsum",
"diff",
"ffill",
"fillna",
"ngroup",
"pct_change",
"rank",
"shift",
]

StataDateFormat: TypeAlias = Literal[
"tc",
Expand Down
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
19 changes: 16 additions & 3 deletions pandas-stubs/core/groupby/generic.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ 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,
Expand All @@ -41,7 +42,6 @@ from pandas._typing import (
ByT,
CorrelationMethod,
Dtype,
GroupByFuncStrs,
IndexLabel,
Level,
ListLike,
Expand Down Expand Up @@ -71,6 +71,15 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
**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],
Expand Down Expand Up @@ -109,7 +118,9 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
**kwargs: Any,
) -> UnknownSeries: ...
@overload
def transform(self, func: GroupByFuncStrs, *args, **kwargs) -> UnknownSeries: ...
def transform(
self, func: TransformReductionListType, *args, **kwargs
) -> UnknownSeries: ...
def filter(
self, func: Callable | str, dropna: bool = ..., *args, **kwargs
) -> Series: ...
Expand Down Expand Up @@ -253,7 +264,9 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
**kwargs: Any,
) -> DataFrame: ...
@overload
def transform(self, func: GroupByFuncStrs, *args, **kwargs) -> DataFrame: ...
def transform(
self, func: TransformReductionListType, *args, **kwargs
) -> DataFrame: ...
def filter(
self, func: Callable, dropna: bool = ..., *args, **kwargs
) -> DataFrame: ...
Expand Down
27 changes: 24 additions & 3 deletions tests/test_series.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1081,7 +1081,7 @@ def test_types_groupby_agg() -> None:

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

check(
assert_type(s.groupby(level=0).agg(sum_sr), "pd.Series[int]"),
Expand Down Expand Up @@ -1119,6 +1119,20 @@ def transform_func(
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:
Expand All @@ -1133,7 +1147,11 @@ def func(s: pd.Series[int]) -> float:
return s.astype(float).min()

s = pd.Series([1, 2, 3, 4])
s.groupby([1, 1, 2, 2]).agg(lambda x: x.astype(float).min())
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,
Expand All @@ -1147,6 +1165,9 @@ def func(s: pd.Series[int]) -> float:
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",
Expand All @@ -1155,7 +1176,7 @@ def func(s: pd.Series[int]) -> float:

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

check(
assert_type(s.groupby(level=0).aggregate(sum_sr), "pd.Series[int]"),
Expand Down