From abfb17ed10c9251d87c87e603ea63c882e8cfb1c Mon Sep 17 00:00:00 2001 From: Quentin FLEURENT NAMBOT Date: Wed, 15 May 2024 16:58:09 +0200 Subject: [PATCH 1/3] fix(frame): add arguments for stack method --- pandas-stubs/core/frame.pyi | 7 ++++++- tests/test_frame.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index f2deecaa8..7331c0d19 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1191,8 +1191,13 @@ class DataFrame(NDFrame, OpsMixin): margins_name: _str = ..., observed: _bool = ..., ) -> DataFrame: ... + @overload + def stack( + self, level: Level | list[Level] = ..., dropna: _bool = ..., sort: _bool = ... + ) -> DataFrame | Series[Any]: ... + @overload def stack( - self, level: Level | list[Level] = ..., dropna: _bool = ... + self, level: Level | list[Level] = ..., future_stack: _bool = ... ) -> DataFrame | Series[Any]: ... def explode( self, column: Sequence[Hashable], ignore_index: _bool = ... diff --git a/tests/test_frame.py b/tests/test_frame.py index dc0040697..dad4439c1 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2256,6 +2256,20 @@ def test_frame_stack() -> None: ), pd.Series, ) + check( + assert_type( + df_multi_level_cols2.stack(0, future_stack=False), + Union[pd.DataFrame, "pd.Series[Any]"], + ), + pd.DataFrame, + ) + check( + assert_type( + df_multi_level_cols2.stack(0, dropna=True, sort=True), + Union[pd.DataFrame, "pd.Series[Any]"], + ), + pd.DataFrame, + ) def test_frame_reindex() -> None: From c41ec61d05d57797aba514166ea375ec4f0cb787 Mon Sep 17 00:00:00 2001 From: Quentin FLEURENT NAMBOT Date: Wed, 15 May 2024 17:29:57 +0200 Subject: [PATCH 2/3] feat: test future_stack only for LTE 2.2 --- tests/test_frame.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index dad4439c1..5cfe4727c 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2256,13 +2256,6 @@ def test_frame_stack() -> None: ), pd.Series, ) - check( - assert_type( - df_multi_level_cols2.stack(0, future_stack=False), - Union[pd.DataFrame, "pd.Series[Any]"], - ), - pd.DataFrame, - ) check( assert_type( df_multi_level_cols2.stack(0, dropna=True, sort=True), @@ -2270,6 +2263,14 @@ def test_frame_stack() -> None: ), pd.DataFrame, ) + if PD_LTE_22: + check( + assert_type( + df_multi_level_cols2.stack(0, future_stack=False), + Union[pd.DataFrame, "pd.Series[Any]"], + ), + pd.DataFrame, + ) def test_frame_reindex() -> None: From dd55298704ef558feafc4efa17f19e662e7f8d79 Mon Sep 17 00:00:00 2001 From: Quentin FLEURENT NAMBOT Date: Wed, 15 May 2024 17:54:49 +0200 Subject: [PATCH 3/3] fix: include both test in PD_LTE_22 condition --- tests/test_frame.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 5cfe4727c..1f4f6bbb4 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2256,13 +2256,6 @@ def test_frame_stack() -> None: ), pd.Series, ) - check( - assert_type( - df_multi_level_cols2.stack(0, dropna=True, sort=True), - Union[pd.DataFrame, "pd.Series[Any]"], - ), - pd.DataFrame, - ) if PD_LTE_22: check( assert_type( @@ -2271,6 +2264,13 @@ def test_frame_stack() -> None: ), pd.DataFrame, ) + check( + assert_type( + df_multi_level_cols2.stack(0, dropna=True, sort=True), + Union[pd.DataFrame, "pd.Series[Any]"], + ), + pd.DataFrame, + ) def test_frame_reindex() -> None: