From adc8df5dc9d0cc098cd784c97acc856cb0e4e089 Mon Sep 17 00:00:00 2001 From: Christopher Yeh Date: Tue, 30 Jul 2024 01:27:37 +0000 Subject: [PATCH 1/3] GroupBy[Series].count() return type should be Series[int] --- pandas-stubs/core/groupby/groupby.pyi | 5 ++++- tests/test_frame.py | 1 + tests/test_series.py | 5 +++++ 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/groupby/groupby.pyi b/pandas-stubs/core/groupby/groupby.pyi index cac030850..75be9578a 100644 --- a/pandas-stubs/core/groupby/groupby.pyi +++ b/pandas-stubs/core/groupby/groupby.pyi @@ -176,7 +176,10 @@ class GroupBy(BaseGroupBy[NDFrameT]): @overload def all(self: GroupBy[DataFrame], skipna: bool = ...) -> DataFrame: ... @final - def count(self) -> NDFrameT: ... + @overload + def count(self: GroupBy[Series]) -> Series[int]: ... + @overload + def count(self: GroupBy[DataFrame]) -> DataFrame: ... @final def mean( self, diff --git a/tests/test_frame.py b/tests/test_frame.py index 8ce0dc356..55219cfa6 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1025,6 +1025,7 @@ def test_types_groupby_methods() -> None: check(assert_type(df.groupby("col1").sum(), pd.DataFrame), pd.DataFrame) check(assert_type(df.groupby("col1").prod(), pd.DataFrame), pd.DataFrame) check(assert_type(df.groupby("col1").sample(), pd.DataFrame), pd.DataFrame) + check(assert_type(df.groupby("col1").count(), pd.DataFrame), pd.DataFrame) check( assert_type(df.groupby("col1").value_counts(normalize=False), "pd.Series[int]"), pd.Series, diff --git a/tests/test_series.py b/tests/test_series.py index 9c99dba41..bfa102a66 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -676,6 +676,11 @@ def test_types_groupby_methods() -> None: check(assert_type(s.groupby(level=0).idxmax(), pd.Series), pd.Series) check(assert_type(s.groupby(level=0).idxmin(), pd.Series), pd.Series) + s2 = pd.Series(["w", "x", "y", "z"], index=[3, 4, 3, 4], dtype=str) + check( + assert_type(s2.groupby(level=0).count(), "pd.Series[int]"), pd.Series, np.int_ + ) + def test_groupby_result() -> None: # GH 142 From b353b02fb73ebdf452fbf538035c7522d8b2ee70 Mon Sep 17 00:00:00 2001 From: Christopher Yeh Date: Wed, 31 Jul 2024 20:53:15 +0000 Subject: [PATCH 2/3] Use np.integer instead of np.int_ --- tests/test_series.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_series.py b/tests/test_series.py index bfa102a66..21f368475 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -678,7 +678,9 @@ def test_types_groupby_methods() -> None: s2 = pd.Series(["w", "x", "y", "z"], index=[3, 4, 3, 4], dtype=str) check( - assert_type(s2.groupby(level=0).count(), "pd.Series[int]"), pd.Series, np.int_ + assert_type(s2.groupby(level=0).count(), "pd.Series[int]"), + pd.Series, + np.integer, ) From f5c2122177b4869e7e1a41061b4106ed0ffb804b Mon Sep 17 00:00:00 2001 From: Christopher Yeh Date: Thu, 1 Aug 2024 22:07:09 +0000 Subject: [PATCH 3/3] Update pyright requirement '>=1.1.369' -> '>=1.1.374' --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 71c22ca55..c5f8f9eb5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ mypy = "1.10.1" pandas = "2.2.2" pyarrow = ">=10.0.1" pytest = ">=7.1.2" -pyright = ">=1.1.369" +pyright = ">= 1.1.374" poethepoet = ">=0.16.5" loguru = ">=0.6.0" typing-extensions = ">=4.4.0"