add proper type when grouping by a Series #708

ClementPinard · 2023-05-24T14:44:44Z

Work in Progress, No test yet

I added a new TypeVar that is the intersection of S1 and ByT.
The type hint for S1 types outside of ByT is not covered anymore, but I suppose that it would theoretically fail in pandas, hence its absence of ByT

Closes Series.groupby(by: Series) should not return SeriesGroupBy[S1, tuple] #706
Tests added: Please use assert_type() to assert the type of any return value

…ies.groupby and DataFrame.groupby

ClementPinard · 2023-05-24T15:07:23Z

Just uploaded tests for iteration over groupby objects.
Note that test are falling now, probably because the typer fails to identify a Series[bool] and thus the correspondig. Would love some guidance over this.

twoertwein · 2023-05-24T15:24:08Z

Would love some guidance over this.

It looks as if it should work :) I would start with adding reveal_type(...) on each line (and break complex statements up) to see which type mypy and pyright infer. The CI stopped after the mypy failure, locally, you can run poe pyright to see whether pyright is happy.

ClementPinard · 2023-05-24T15:26:05Z

My vscode was happy so I guess pyright was happy.

However, mypy do fail on local too.

I'll do some additional investigation.

ClementPinard · 2023-05-24T15:42:50Z

Interesting take with reveal_types :

def test_types_groupby_iter() -> None:
    s = pd.Series([1, 1, 2])
    series_groupby: pd.Series[bool] = pd.Series([True, True, False])
    first_group = next(iter(s.groupby(series_groupby)))
    reveal_type(s.groupby(series_groupby))
    reveal_type(s.groupby(series_groupby).__iter__())
    reveal_type(iter(s.groupby(series_groupby)))

Outputs

tests/test_series.py:738: note: Revealed type is "pandas.core.groupby.generic.SeriesGroupBy[Any, builtins.bool]"
tests/test_series.py:739: note: Revealed type is "typing.Iterator[Tuple[builtins.bool, pandas.core.series.Series[Any]]]"
tests/test_series.py:740: note: Revealed type is "typing.Iterator[Any]"

How can the type of iter(something) be different from something.__iter__ ? Is there some double check I should do ?

Even more puzzling is that the type is right for DataFrames :

def test_types_groupby_iter() -> None:
    df = pd.DataFrame(data={"col1": [1, 1, 2], "col2": [3, 4, 5]})
    series_groupby: pd.Series[bool] = pd.Series([True, True, False])
    first_group = next(iter(df.groupby(series_groupby)))
    reveal_type(df.groupby(series_groupby))
    reveal_type(df.groupby(series_groupby).__iter__())
    reveal_type(iter(df.groupby(series_groupby)))

output

tests/test_frame.py:992: note: Revealed type is "pandas.core.groupby.generic.DataFrameGroupBy[builtins.bool]"
tests/test_frame.py:993: note: Revealed type is "typing.Iterator[Tuple[builtins.bool, pandas.core.frame.DataFrame]]"
tests/test_frame.py:994: note: Revealed type is "typing.Iterator[Tuple[builtins.bool, pandas.core.frame.DataFrame]]"

But I don't see any particular additional type hint DataFrameGroupBy regarding iteration compared to SeriesGroupBy ...

Dr-Irv · 2023-05-24T20:11:22Z

How can the type of iter(something) be different from something.__iter__ ? Is there some double check I should do ?

I think the issue is that when you are doing s.groupby(series_groupby), s is Series[Any] because there was no dtype information supplied. If you had created s=pd.Series([1, 1, 2], dtype=int), then you'd get the results you expect.

NOTE - my ability to help out will be pretty limited between now (5/24) and 6/4 due to vacation.

ClementPinard · 2023-05-24T20:24:59Z

You are right, mypy now works !

Did not know the dtype was useful to help the type checker.

twoertwein · 2023-05-24T21:02:51Z

NOTE - my ability to help out will be pretty limited between now (5/24) and 6/4 due to vacation.

I will merge low-risk PRs during that timeframe.

pandas-stubs/_typing.pyi

gandhis1 · 2023-05-24T23:12:03Z

How can the type of iter(something) be different from something.__iter__ ? Is there some double check I should do ?

I think the issue is that when you are doing s.groupby(series_groupby), s is Series[Any] because there was no dtype information supplied. If you had created s=pd.Series([1, 1, 2], dtype=int), then you'd get the results you expect.

NOTE - my ability to help out will be pretty limited between now (5/24) and 6/4 due to vacation.

Maybe this is deserving of a separate issue, but shouldn't we be able to create an initializer overload that takes a list[T] and returns a pd.Series[T]?

Dr-Irv · 2023-05-25T13:47:13Z

Maybe this is deserving of a separate issue, but shouldn't we be able to create an initializer overload that takes a list[T] and returns a pd.Series[T]?

Can you create a separate issue? I think this might work.

Dr-Irv

A couple small things from me, plus address issue to add comment about SeriesByT

Dr-Irv · 2023-05-25T18:27:37Z

tests/test_frame.py

+        assert_type(first_group[0], "bool"),
+        bool,
+    )
+    check(
+        assert_type(first_group[1], "pd.DataFrame"),


You don't need quotes around the types in the two assert_type() statements, because those are valid types at runtime.

Dr-Irv · 2023-05-25T18:27:50Z

tests/test_series.py

+    series_groupby = pd.Series([True, True, False], dtype=bool)
+    first_group = next(iter(s.groupby(series_groupby)))
+    check(
+        assert_type(first_group[0], "bool"),


no quotes needed here

Dr-Irv · 2023-05-25T18:28:18Z

tests/test_series.py

+        assert_type(first_group[1], "pd.Series[int]"),
+        pd.Series,


Suggested change

assert_type(first_group[1], "pd.Series[int]"),

pd.Series,

assert_type(first_group[1], "pd.Series[int]"),

pd.Series, np.integer

You do need quotes here because pd.Series[int] is not a type at runtime.

ClementPinard · 2023-05-30T21:39:33Z

Thank you for the review, I do believe I have addressed the comments

Dr-Irv

Because of a recent change to the definition of S1, can you merge with main, address the comment below, and if the tests pass, we should be good to go.

Dr-Irv · 2023-06-01T01:40:41Z

pandas-stubs/_typing.pyi

+# Essentially, an intersection between Series S1 TypeVar, and ByT TypeVar
+SeriesByT = TypeVar(
+    "SeriesByT",
+    str,


We've been recently advised to no longer use constrained TypeVars, so can you change this to SeriesByT = TypeVar("SeriesByT", bound=str | bytes | ...) We made the change for S1 in a recent PR

Hi,

What is the difference between TypeVar("SeriesByT", bound = str | bytes ...) and TypeVar("SeriesByT, str, byts, ...) ?

I tried to replace the SeriesByT with the bound keyword, but then the mypy did not pass anymore :

Poe => mypy =========================================== Beginning: 'Run mypy on 'tests' (using the local stubs) and on the local stubs' =========================================== pandas-stubs/core/series.pyi:648: error: Type variable "SeriesByT" not valid as type argument value for "SeriesGroupBy" [type-var] pandas-stubs/core/frame.pyi:1100: error: Type variable "SeriesByT" not valid as type argument value for "DataFrameGroupBy" [type-var] Found 2 errors in 2 files (checked 224 source files) =========================================== Step: 'Run mypy on 'tests' (using the local stubs) and on the local stubs' failed! ===========================================

bound allows the type to be a union of all the listed types. Without it, the type has to be one of the listed (sub-)types.

You probably also need to adjust the type used for SeriesGroupBy and DataFrameGroupBy to use bound.

Thanks I actually tried that and got rid of the errors.

Dr-Irv

thanks @ClementPinard

ClementPinard added 2 commits May 24, 2023 16:41

remove Series from GroupByObjectNonScalar, add new orverloads for Ser…

ac92a2f

…ies.groupby and DataFrame.groupby

Add tests for iteration over groupby

413f08d

Add dtype in tests

4eeb52a

twoertwein reviewed May 24, 2023

View reviewed changes

pandas-stubs/_typing.pyi Show resolved Hide resolved

gandhis1 mentioned this pull request May 25, 2023

Create pd.Series initializer overload to propagate the generic type of an input container #710

Closed

Dr-Irv requested changes May 25, 2023

View reviewed changes

address PR comments

b984743

twoertwein approved these changes May 31, 2023

View reviewed changes

Dr-Irv requested changes Jun 1, 2023

View reviewed changes

ClementPinard and others added 2 commits June 1, 2023 16:12

Merge branch 'pandas-dev:main' into feature/groupby_series

903b662

Use bound method for ByT and SeriesByT typevars

8d692fb

Dr-Irv approved these changes Jun 2, 2023

View reviewed changes

Dr-Irv merged commit d873a46 into pandas-dev:main Jun 2, 2023

ClementPinard deleted the feature/groupby_series branch June 2, 2023 12:47

Uh oh!

add proper type when grouping by a Series #708

add proper type when grouping by a Series #708

Uh oh!

Conversation

ClementPinard commented May 24, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ClementPinard commented May 24, 2023

Uh oh!

twoertwein commented May 24, 2023

Uh oh!

ClementPinard commented May 24, 2023

Uh oh!

ClementPinard commented May 24, 2023

Uh oh!

Dr-Irv commented May 24, 2023

Uh oh!

ClementPinard commented May 24, 2023

Uh oh!

twoertwein commented May 24, 2023

Uh oh!

Uh oh!

gandhis1 commented May 24, 2023

Uh oh!

Dr-Irv commented May 25, 2023

Uh oh!

Dr-Irv left a comment

Choose a reason for hiding this comment

Uh oh!

Dr-Irv May 25, 2023

Choose a reason for hiding this comment

Uh oh!

Dr-Irv May 25, 2023

Choose a reason for hiding this comment

Uh oh!

Dr-Irv May 25, 2023

Choose a reason for hiding this comment

Uh oh!

ClementPinard commented May 30, 2023

Uh oh!

Dr-Irv left a comment

Choose a reason for hiding this comment

Uh oh!

Dr-Irv Jun 1, 2023

Choose a reason for hiding this comment

Uh oh!

ClementPinard Jun 1, 2023

Choose a reason for hiding this comment

Uh oh!

twoertwein Jun 1, 2023

Choose a reason for hiding this comment

Uh oh!

ClementPinard Jun 1, 2023

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!

ClementPinard commented May 24, 2023 •

edited

Loading