GH456 First attempt GroupBy.transform improved typing #1242

loicdiridollou · 2025-06-06T21:16:56Z

Closes Create more precise annotations for DataFrame/Series GroupBy operations (agg, apply, transform) #456
Tests added: Please use assert_type() to assert the type of any return value

pandas-stubs/_typing.pyi

pandas-stubs/core/groupby/generic.pyi

Dr-Irv

Good idea. I think it won't close 456, because that was also about agg .

pandas-stubs/core/groupby/generic.pyi

tests/test_series.py

pandas-stubs/_typing.pyi

loicdiridollou · 2025-06-06T23:39:22Z

For

I think you should just copy what is in the source pandas.core.groupby.base.py into pandas.core.groupby.base.pyi so they are equivalent, but change the frozenset to be Literal

is there a convention to name the TypeAlias or should we keep the same name as in the pandas code?

Dr-Irv · 2025-06-07T14:37:07Z

For
I think you should just copy what is in the source pandas.core.groupby.base.py into pandas.core.groupby.base.pyi so they are equivalent, but change the frozenset to be Literal
is there a convention to name the TypeAlias or should we keep the same name as in the pandas code?

Probably shouldn't use the same name, because then the types won't match. Not that anyone should be importing that anyway. So keep the same name and append _type after the name.

Dr-Irv · 2025-06-07T14:38:05Z

pandas-stubs/core/groupby/base.pyi


 @dataclasses.dataclass(order=True, frozen=True)
 class OutputKey:
    label: Hashable
    position: int
+
+reduction_kernels: TypeAlias = Literal[


as commented elsewhere, lets call this reduction_kernels_type so that the types of the objects line up.

Dr-Irv · 2025-06-07T14:39:15Z

tests/test_series.py

+            # type of `sum` not well inferred by mypy
+            return sum(s)


why not use s.sum() ?

The issue was if I passed sum the builtin one in the aggregate method directly, here is does not change anything.

Dr-Irv · 2025-06-07T14:40:24Z

tests/test_series.py

+        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())


don't you want a check(assert_type(... here?

Correct my mistake

Turns out the inference on the fly of lambdas is not super clear so you need to define the function on the side to have the right types.

Yes, that is an issue with lambda functions.

Actually, I think you can have a test of

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

which would be worthwhile

Dr-Irv · 2025-06-07T14:40:43Z

tests/test_series.py

+            # type of `sum` not well inferred by mypy
+            return sum(s)


use s.sum()

Dr-Irv · 2025-06-07T15:32:53Z

tests/test_series.py

        check(
            assert_type(s.groupby(level=0).agg([min, sum]), pd.DataFrame), pd.DataFrame
        )


+def test_types_groupby_transform() -> None:


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

Dr-Irv · 2025-06-07T17:07:57Z

pandas-stubs/core/groupby/generic.pyi

    def aggregate(
        self,
        func: AggFuncTypeBase | None = ...,
+        /,
        *args,


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.

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

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?

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)

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.

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

Dr-Irv · 2025-06-10T13:25:37Z

tests/test_series.py

@@ -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]
+    # check(assert_type(s.groupby([1,1,2,2]).agg(lambda x: x.astype(float).min()), pd.Series), pd.Series, float)


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].

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?

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.

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

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.

Dr-Irv

Just have to adjust that comment, then we are good to go.

Dr-Irv · 2025-06-11T16:59:05Z

tests/test_series.py

+        np.floating,
+    )
+
+    # test below passes with mypy but pyright correctly sees it as pd.Series[float]


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

Dr-Irv

thanks @loicdiridollou

GH456 First attempt GroupBy.transform improved typing

020f93d

loicdiridollou commented Jun 6, 2025

View reviewed changes

pandas-stubs/_typing.pyi Outdated Show resolved Hide resolved

loicdiridollou commented Jun 6, 2025

View reviewed changes

pandas-stubs/core/groupby/generic.pyi Show resolved Hide resolved

Dr-Irv requested changes Jun 6, 2025

View reviewed changes

loicdiridollou added 2 commits June 6, 2025 19:23

GH456 Attempt GroupBy.aggregate improved typing

106a6f5

GH456 Attempt GroupBy.aggregate improved typing

3bba101

Dr-Irv requested changes Jun 7, 2025

View reviewed changes

GH456 PR Feedback

053b7e7

loicdiridollou requested a review from Dr-Irv June 7, 2025 15:11

Dr-Irv requested changes Jun 7, 2025

View reviewed changes

Dr-Irv reviewed Jun 7, 2025

View reviewed changes

loicdiridollou added 3 commits June 7, 2025 16:16

GH456 PR Feedback

4141a06

GH456 PR Feedback

f9863d0

GH456 PR Feedback

e26b4c1

Dr-Irv requested changes Jun 10, 2025

View reviewed changes

Dr-Irv requested changes Jun 11, 2025

View reviewed changes

GH456 PR Feedback

96abf3b

loicdiridollou requested a review from Dr-Irv June 13, 2025 22:21

Dr-Irv approved these changes Jun 13, 2025

View reviewed changes

Dr-Irv merged commit b12c28d into pandas-dev:main Jun 13, 2025
13 checks passed

Uh oh!

GH456 First attempt GroupBy.transform improved typing #1242

GH456 First attempt GroupBy.transform improved typing #1242

Uh oh!

Conversation

loicdiridollou commented Jun 6, 2025

Uh oh!

Uh oh!

Uh oh!

Dr-Irv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

loicdiridollou commented Jun 6, 2025

Uh oh!

Dr-Irv commented Jun 7, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Dr-Irv left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Dr-Irv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!