From e13c91e7c1aaff914949c9430d75d07b28dd0cad Mon Sep 17 00:00:00 2001 From: Aftab Uddin Date: Wed, 20 Mar 2024 05:29:25 +0600 Subject: [PATCH 1/4] #885 Added Sequence[_str] to the values argument in the stub --- pandas-stubs/core/frame.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index fe3fef2f7..7ce2f5d61 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1125,7 +1125,7 @@ class DataFrame(NDFrame, OpsMixin): ) -> DataFrame: ... def pivot_table( self, - values: _str | None = ..., + values: _str | None | Sequence[_str] = ..., index: _str | Grouper | Sequence | None = ..., columns: _str | Grouper | Sequence | None = ..., aggfunc=..., From 1a17dfd6e99e72b867170c490ceba48ac7ff738f Mon Sep 17 00:00:00 2001 From: Aftab Uddin Date: Thu, 21 Mar 2024 01:11:18 +0600 Subject: [PATCH 2/4] Added the test case test_types_pivot_table() for the changed TYP of values argument in pivot_table --- tests/test_frame.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/test_frame.py b/tests/test_frame.py index 0eab81915..28003ad62 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -932,6 +932,23 @@ def test_types_pivot() -> None: ) +def test_types_pivot_table() -> None: + df = pd.DataFrame( + data={ + "col1": ["first", "second", "third", "fourth"], + "col2": [50, 70, 56, 111], + "col3": ["A", "B", "C", "D"], + "col4": [100, 102, 500, 600], + } + ) + check( + assert_type( + df.pivot_table(index="col2", columns="col4", values=["col1", "col3"]), pd.DataFrame, + ), + pd.DataFrame, + ) + + def test_types_groupby() -> None: df = pd.DataFrame(data={"col1": [1, 1, 2], "col2": [3, 4, 5], "col3": [0, 1, 0]}) df.index.name = "ind" From d2df21f6ea2da200371ffc20ee61b8d74653cd50 Mon Sep 17 00:00:00 2001 From: Aftab Uddin Date: Thu, 21 Mar 2024 04:14:56 +0600 Subject: [PATCH 3/4] Changed the argument values for the test case test_types_pivot_table --- tests/test_frame.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 28003ad62..8e5957b9f 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -943,7 +943,7 @@ def test_types_pivot_table() -> None: ) check( assert_type( - df.pivot_table(index="col2", columns="col4", values=["col1", "col3"]), pd.DataFrame, + df.pivot_table(index="col1", columns="col3", values=["col2", "col4"]), pd.DataFrame, ), pd.DataFrame, ) From d9f0ea90134a5e114a4f1b53b4a4d4b0f6da8c15 Mon Sep 17 00:00:00 2001 From: Aftab Uddin Date: Thu, 21 Mar 2024 05:16:27 +0600 Subject: [PATCH 4/4] Fixed formatting with black --- tests/test_frame.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 8e5957b9f..af3183b93 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -943,11 +943,12 @@ def test_types_pivot_table() -> None: ) check( assert_type( - df.pivot_table(index="col1", columns="col3", values=["col2", "col4"]), pd.DataFrame, + df.pivot_table(index="col1", columns="col3", values=["col2", "col4"]), + pd.DataFrame, ), pd.DataFrame, ) - + def test_types_groupby() -> None: df = pd.DataFrame(data={"col1": [1, 1, 2], "col2": [3, 4, 5], "col3": [0, 1, 0]})