From 8a41508683212d6e861f07d053d86df796460124 Mon Sep 17 00:00:00 2001 From: Quentin FLEURENT NAMBOT Date: Wed, 15 May 2024 18:33:03 +0200 Subject: [PATCH] fix(frame): add arguments for stack method (#926) * fix(frame): add arguments for stack method * feat: test future_stack only for LTE 2.2 * fix: include both test in PD_LTE_22 condition --------- Co-authored-by: Quentin FLEURENT NAMBOT --- pandas-stubs/core/frame.pyi | 7 ++++++- tests/test_frame.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index f2deecaa..7331c0d1 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 dc004069..1f4f6bbb 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2256,6 +2256,21 @@ def test_frame_stack() -> None: ), pd.Series, ) + if PD_LTE_22: + 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: