diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index fe3fef2f7..7ce2f5d61 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1125,7 +1125,7 @@ class DataFrame(NDFrame, OpsMixin): ) -> DataFrame: ... def pivot_table( self, - values: _str | None = ..., + values: _str | None | Sequence[_str] = ..., index: _str | Grouper | Sequence | None = ..., columns: _str | Grouper | Sequence | None = ..., aggfunc=..., diff --git a/tests/test_frame.py b/tests/test_frame.py index 0eab81915..af3183b93 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -932,6 +932,24 @@ def test_types_pivot() -> None: ) +def test_types_pivot_table() -> None: + df = pd.DataFrame( + data={ + "col1": ["first", "second", "third", "fourth"], + "col2": [50, 70, 56, 111], + "col3": ["A", "B", "C", "D"], + "col4": [100, 102, 500, 600], + } + ) + check( + assert_type( + df.pivot_table(index="col1", columns="col3", values=["col2", "col4"]), + pd.DataFrame, + ), + pd.DataFrame, + ) + + def test_types_groupby() -> None: df = pd.DataFrame(data={"col1": [1, 1, 2], "col2": [3, 4, 5], "col3": [0, 1, 0]}) df.index.name = "ind"