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: