From c4d24b23a8d3e3f1e0cb354fed16fb1435b46e9d Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Fri, 2 Dec 2022 01:56:41 +0530 Subject: [PATCH 1/4] ty --- pandas-stubs/core/groupby/generic.pyi | 1 + tests/test_frame.py | 11 +++++++++++ 2 files changed, 12 insertions(+) diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index c4a56f8b0..ff8ba9e67 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -359,3 +359,4 @@ class DataFrameGroupBy(GroupBy): ascending: bool = ..., dropna: bool = ..., ) -> Series[float]: ... + def __getattr__(self, name: str) -> Series: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 44b5317cb..9fd92b026 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2097,3 +2097,14 @@ def test_groupby_and_transform() -> None: check(assert_type(c2, pd.DataFrame), pd.DataFrame) check(assert_type(c3, pd.Series), pd.Series) check(assert_type(c4, pd.Series), pd.Series) + + +def test_getattr_and_dataframe_groupby() -> None: + df = pd.DataFrame( + { + 'C': [1, 5, 5, 2, 5, 5], + 'D': [2.0, 5.0, 8.0, 1.0, 2.0, 9.0], + } + ) + gb = df.groupby('C').D + check(assert_type(gb, pd.Series), pd.Series) From 00747b996c2fe5c8825d8871e5b984d3fcdb4591 Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Mon, 5 Dec 2022 21:40:13 +0530 Subject: [PATCH 2/4] changed the return value to SeriesGroupBy and corrected the tests likewise --- pandas-stubs/core/groupby/generic.pyi | 2 +- tests/test_frame.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index ff8ba9e67..b1890d522 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -359,4 +359,4 @@ class DataFrameGroupBy(GroupBy): ascending: bool = ..., dropna: bool = ..., ) -> Series[float]: ... - def __getattr__(self, name: str) -> Series: ... + def __getattr__(self, name: str) -> SeriesGroupBy: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 9fd92b026..8994deb88 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -30,6 +30,7 @@ ensure_clean, getSeriesData, ) +from pandas.core.groupby.generic import SeriesGroupBy from pandas.core.resample import Resampler # noqa: F401 import pytest from typing_extensions import assert_type @@ -2102,9 +2103,9 @@ def test_groupby_and_transform() -> None: def test_getattr_and_dataframe_groupby() -> None: df = pd.DataFrame( { - 'C': [1, 5, 5, 2, 5, 5], - 'D': [2.0, 5.0, 8.0, 1.0, 2.0, 9.0], + "C": [1, 5, 5, 2, 5, 5], + "D": [2.0, 5.0, 8.0, 1.0, 2.0, 9.0], } ) - gb = df.groupby('C').D - check(assert_type(gb, pd.Series), pd.Series) + gb = df.groupby("C").D + check(assert_type(gb, SeriesGroupBy), SeriesGroupBy) From db487f8faa970ecae9fa2f538061103a4b2b750c Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Tue, 6 Dec 2022 22:37:27 +0530 Subject: [PATCH 3/4] changed test to test '__getattr()' indirectly --- pandas-stubs/core/groupby/generic.pyi | 2 +- tests/test_frame.py | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index b1890d522..ff8ba9e67 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -359,4 +359,4 @@ class DataFrameGroupBy(GroupBy): ascending: bool = ..., dropna: bool = ..., ) -> Series[float]: ... - def __getattr__(self, name: str) -> SeriesGroupBy: ... + def __getattr__(self, name: str) -> Series: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 8994deb88..15568fcd3 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -30,7 +30,6 @@ ensure_clean, getSeriesData, ) -from pandas.core.groupby.generic import SeriesGroupBy from pandas.core.resample import Resampler # noqa: F401 import pytest from typing_extensions import assert_type @@ -2103,9 +2102,9 @@ def test_groupby_and_transform() -> None: def test_getattr_and_dataframe_groupby() -> None: df = pd.DataFrame( { - "C": [1, 5, 5, 2, 5, 5], - "D": [2.0, 5.0, 8.0, 1.0, 2.0, 9.0], + "col1": [1, 5, 5, 2, 5, 5], + "col2": [2.0, 5.0, 8.0, 1.0, 2.0, 9.0], + "col3": [9,8,7,5,6,1] } ) - gb = df.groupby("C").D - check(assert_type(gb, SeriesGroupBy), SeriesGroupBy) + check(assert_type(df.groupby("col1").__getattr__("col3"), pd.Series), pd.Series) From 40da7e92f72bd587380683dae20945cd77b367dc Mon Sep 17 00:00:00 2001 From: ramvikrams Date: Wed, 7 Dec 2022 10:02:16 +0530 Subject: [PATCH 4/4] modified the tests so SeriesGroupBy is not tested directly but indirectly --- pandas-stubs/core/groupby/generic.pyi | 2 +- tests/test_frame.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index ff8ba9e67..b1890d522 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -359,4 +359,4 @@ class DataFrameGroupBy(GroupBy): ascending: bool = ..., dropna: bool = ..., ) -> Series[float]: ... - def __getattr__(self, name: str) -> Series: ... + def __getattr__(self, name: str) -> SeriesGroupBy: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 15568fcd3..a760a679d 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2101,10 +2101,10 @@ def test_groupby_and_transform() -> None: def test_getattr_and_dataframe_groupby() -> None: df = pd.DataFrame( - { - "col1": [1, 5, 5, 2, 5, 5], - "col2": [2.0, 5.0, 8.0, 1.0, 2.0, 9.0], - "col3": [9,8,7,5,6,1] - } + data={"col1": [1, 1, 2], "col2": [3, 4, 5], "col3": [0, 1, 0], 0: [-1, -1, -1]} + ) + check(assert_type(df.groupby("col1").col3.agg(min), pd.Series), pd.Series) + check( + assert_type(df.groupby("col1").col3.agg([min, max]), pd.DataFrame), + pd.DataFrame, ) - check(assert_type(df.groupby("col1").__getattr__("col3"), pd.Series), pd.Series)