diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 2a9875d83..90560dc2a 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 22ef5c762..1982d54bc 100644 --- a/pandas-stubs/core/series.pyi +++ b/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 diff --git a/tests/test_frame.py b/tests/test_frame.py index a951c8aa0..d58bd7f89 100644 --- a/tests/test_frame.py +++ b/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 diff --git a/tests/test_series.py b/tests/test_series.py index a2ace21e5..3f25d981e 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -3902,7 +3902,6 @@ 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( @@ -3910,3 +3909,16 @@ def test_cumsum_timedelta() -> None: 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, + )