update types, add tests

Author: Stephen Randall

pandas-stubs/core/frame.pyi

@@ -1317,18 +1317,18 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     ) -> Self: ...
     @overload
     def stack(
-        self, level: Level | list[Level] = ..., dropna: _bool = ..., sort: _bool = ...
+        self, level: IndexLabel = ..., dropna: _bool = ..., sort: _bool = ...
     ) -> Self | Series: ...
     @overload
     def stack(
-        self, level: Level | list[Level] = ..., future_stack: _bool = ...
+        self, level: IndexLabel = ..., future_stack: _bool = ...
     ) -> Self | Series: ...
     def explode(
         self, column: Sequence[Hashable], ignore_index: _bool = ...
     ) -> Self: ...
     def unstack(
         self,
-        level: Level = ...,
+        level: IndexLabel = ...,
         fill_value: Scalar | None = ...,
         sort: _bool = ...,
     ) -> Self | Series: ...

pandas-stubs/core/series.pyi

@@ -934,7 +934,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     def explode(self) -> Series[S1]: ...
     def unstack(
         self,
-        level: Level = ...,
+        level: IndexLabel = ...,
         fill_value: int | _str | dict | None = ...,
     ) -> DataFrame: ...
     @overload

tests/test_frame.py

@@ -3124,6 +3124,12 @@ def test_frame_stack() -> None:
             assert_type(df_multi_level_cols2.stack(0), Union[pd.DataFrame, pd.Series]),
             pd.DataFrame,
         )
+        check(
+            assert_type(
+                df_multi_level_cols2.stack([0]), Union[pd.DataFrame, pd.Series]
+            ),
+            pd.DataFrame,
+        )
         check(
             assert_type(
                 df_multi_level_cols2.stack([0, 1]),
@@ -4461,6 +4467,7 @@ def test_unstack() -> None:
     ).set_index([0, 1])
 
     check(assert_type(df.unstack(0), pd.DataFrame | pd.Series), pd.DataFrame)
+    check(assert_type(df.unstack([0]), pd.DataFrame | pd.Series), pd.DataFrame)
     check(
         assert_type(
             df.unstack(1, fill_value=pd.Timestamp(2023, 4, 5)), pd.DataFrame | pd.Series

tests/test_series.py

@@ -3902,11 +3902,22 @@ def test_series_items() -> None:
 
 
 def test_cumsum_timedelta() -> None:
-
     s = pd.Series(pd.to_timedelta([1, 2, 3], "h"))
     check(assert_type(s.cumsum(), "TimedeltaSeries"), pd.Series, pd.Timedelta)
     check(
         assert_type(pd.Timestamp(0) + s.cumsum(), "TimestampSeries"),
         pd.Series,
         pd.Timestamp,
     )
+
+
+def test_series_unstack() -> None:
+    df = pd.DataFrame([[1, 3, 5], [2, 4, 6]])
+    s = df.transpose().stack([*range(df.index.nlevels)])
+    check(
+        assert_type(
+            s.unstack([*range(s.index.nlevels // 2)]),
+            Union[pd.Series, pd.DataFrame],
+        ),
+        pd.DataFrame,
+    )