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 <[email protected]>

Author: rinzool

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 = ...

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: