Update types for Series.unstack, DataFrame.stack, and DataFrame.unstack

[stephenlrandall]

• Closes Incorrect type for Series.unstack #1256
• Tests added: One for Series.unstack, and modified the test for DataFrame.unstack

[Dr-Irv]

Can you do check(assert_type(s, pd.Series), pd.Series, np.int64) here? (I think that's the right type)

[stephenlrandall]

That fails with:

tests/test_series.py:3917: error: Expression is of type "DataFrame | Series[Any]", not "Series[Any]"  [assert-type]
Found 1 error in 1 file (checked 224 source files)

What I can do instead is:

def test_series_unstack() -> None:
    df = pd.DataFrame([[1, 3, 5], [2, 4, 6]])
    s = cast(pd.Series, df.transpose().stack([*range(df.index.nlevels)]))
    check(assert_type(s, pd.Series), pd.Series, np.int64)
    check(
        assert_type(
            s.unstack([*range(s.index.nlevels // 2)]),
            pd.DataFrame,
        ),
        pd.DataFrame,
    )

where I explicitly cast the stacked object to a Series. Or would you rather have something else there?

[Dr-Irv]

Replace s = cast(pd.Series, df.transpose().stack([*range(df.index.nlevels)])) with

s = df.transpose().stack([*range(df.index.nlevels)]))
check(assert_type(s, Union[pd.Series, pd.DataFrame), pd.Series)

This then says from a typing perspective, you could get the Union, but at runtime, you will get a Series

[stephenlrandall]

Makes sense -- updated the test.

[Dr-Irv]

thanks @stephenlrandall