From 06200246463c0fca8630b26f8f0e685d095fe741 Mon Sep 17 00:00:00 2001
From: Mudassir <muddi900@gmail.com>
Date: Sat, 31 Dec 2022 16:53:43 +0500
Subject: [PATCH] added list[Levels] as per #489 and added appropriate tests

added a generic Any for series output

created tests for stack

wrapping generics in quotes for compatibility

style update
---
 pandas-stubs/core/frame.pyi |  4 +++-
 tests/test_frame.py         | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi
index 4eadbf022..a76f3d83a 100644
--- a/pandas-stubs/core/frame.pyi
+++ b/pandas-stubs/core/frame.pyi
@@ -997,7 +997,9 @@ class DataFrame(NDFrame, OpsMixin):
         margins_name: _str = ...,
         observed: _bool = ...,
     ) -> DataFrame: ...
-    def stack(self, level: Level = ..., dropna: _bool = ...) -> DataFrame | Series: ...
+    def stack(
+        self, level: Level | list[Level] = ..., dropna: _bool = ...
+    ) -> DataFrame | Series[Any]: ...
     def explode(self, column: _str | tuple, ignore_index: _bool = ...) -> DataFrame: ...
     def unstack(
         self,
diff --git a/tests/test_frame.py b/tests/test_frame.py
index 02f45bc1f..4a34be2ef 100644
--- a/tests/test_frame.py
+++ b/tests/test_frame.py
@@ -1798,6 +1798,27 @@ def test_frame_index_numpy() -> None:
     pd.DataFrame([[1.0, 2.0], [3.0, 4.0]], columns=["a", "b"], index=i)
 
 
+def test_frame_stack() -> None:
+
+    multicol2 = pd.MultiIndex.from_tuples([("weight", "kg"), ("height", "m")])
+    df_multi_level_cols2 = pd.DataFrame(
+        [[1.0, 2.0], [3.0, 4.0]], index=["cat", "dog"], columns=multicol2
+    )
+
+    check(
+        assert_type(
+            df_multi_level_cols2.stack(0), Union[pd.DataFrame, "pd.Series[Any]"]
+        ),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            df_multi_level_cols2.stack([0, 1]), Union[pd.DataFrame, "pd.Series[Any]"]
+        ),
+        pd.Series,
+    )
+
+
 def test_frame_reindex() -> None:
     # GH 84
     df = pd.DataFrame({"a": [1, 2, 3]}, index=[0, 1, 2])