Update types for Series.unstack, DataFrame.stack, and DataFrame.unstack (#1259)

* update types, add tests

* remove unnecessary test

* added check

---------

Co-authored-by: Stephen Randall <[email protected]>

Author: stephenlrandall

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

@@ -4461,6 +4461,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,23 @@ 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, Union[pd.Series, pd.DataFrame]), pd.Series)
+    check(
+        assert_type(
+            s.unstack([*range(s.index.nlevels // 2)]),
+            Union[pd.Series, pd.DataFrame],
+        ),
+        pd.DataFrame,
+    )