fix: #1212 use typing from pandas.core.reshape.pivot

Author: cmp0xff

pandas-stubs/core/frame.pyi

@@ -32,7 +32,6 @@ from pandas import (
 from pandas.core.arraylike import OpsMixin
 from pandas.core.generic import NDFrame
 from pandas.core.groupby.generic import DataFrameGroupBy
-from pandas.core.groupby.grouper import Grouper
 from pandas.core.indexers import BaseIndexer
 from pandas.core.indexes.base import (
     Index,
@@ -50,6 +49,11 @@ from pandas.core.indexing import (
     _LocIndexer,
 )
 from pandas.core.interchange.dataframe_protocol import DataFrame as DataFrameXchg
+from pandas.core.reshape.pivot import (
+    PivotTableColumnsTypes,
+    PivotTableIndexTypes,
+    PivotTableValuesTypes,
+)
 from pandas.core.series import Series
 from pandas.core.window import (
     Expanding,
@@ -1287,9 +1291,9 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     ) -> Self: ...
     def pivot_table(
         self,
-        values: _str | None | Sequence[_str] = ...,
-        index: _str | Grouper | Sequence | None = ...,
-        columns: _str | Grouper | Sequence | None = ...,
+        values: PivotTableValuesTypes = ...,
+        index: PivotTableIndexTypes = ...,
+        columns: PivotTableColumnsTypes = ...,
         aggfunc=...,
         fill_value: Scalar | None = ...,
         margins: _bool = ...,

pandas-stubs/core/reshape/pivot.pyi

@@ -51,19 +51,22 @@ _NonIterableHashable: TypeAlias = (
     | pd.Timedelta
 )
 
-_PivotTableIndexTypes: TypeAlias = Label | list[HashableT1] | Series | Grouper | None
-_PivotTableColumnsTypes: TypeAlias = Label | list[HashableT2] | Series | Grouper | None
+PivotTableIndexTypes: TypeAlias = Label | Sequence[HashableT1] | Series | Grouper | None
+PivotTableColumnsTypes: TypeAlias = (
+    Label | Sequence[HashableT2] | Series | Grouper | None
+)
+PivotTableValuesTypes: TypeAlias = Label | Sequence[HashableT3] | None
 
 _ExtendedAnyArrayLike: TypeAlias = AnyArrayLike | ArrayLike
 
 @overload
 def pivot_table(
     data: DataFrame,
-    values: Label | list[HashableT3] | None = ...,
-    index: _PivotTableIndexTypes = ...,
-    columns: _PivotTableColumnsTypes = ...,
+    values: PivotTableValuesTypes = ...,
+    index: PivotTableIndexTypes = ...,
+    columns: PivotTableColumnsTypes = ...,
     aggfunc: (
-        _PivotAggFunc | list[_PivotAggFunc] | Mapping[Hashable, _PivotAggFunc]
+        _PivotAggFunc | Sequence[_PivotAggFunc] | Mapping[Hashable, _PivotAggFunc]
     ) = ...,
     fill_value: Scalar | None = ...,
     margins: bool = ...,
@@ -77,12 +80,12 @@ def pivot_table(
 @overload
 def pivot_table(
     data: DataFrame,
-    values: Label | list[HashableT3] | None = ...,
+    values: PivotTableValuesTypes = ...,
     *,
     index: Grouper,
-    columns: _PivotTableColumnsTypes | Index | npt.NDArray = ...,
+    columns: PivotTableColumnsTypes | Index | npt.NDArray = ...,
     aggfunc: (
-        _PivotAggFunc | list[_PivotAggFunc] | Mapping[Hashable, _PivotAggFunc]
+        _PivotAggFunc | Sequence[_PivotAggFunc] | Mapping[Hashable, _PivotAggFunc]
     ) = ...,
     fill_value: Scalar | None = ...,
     margins: bool = ...,
@@ -94,12 +97,12 @@ def pivot_table(
 @overload
 def pivot_table(
     data: DataFrame,
-    values: Label | list[HashableT3] | None = ...,
-    index: _PivotTableIndexTypes | Index | npt.NDArray = ...,
+    values: PivotTableValuesTypes = ...,
+    index: PivotTableIndexTypes | Index | npt.NDArray = ...,
     *,
     columns: Grouper,
     aggfunc: (
-        _PivotAggFunc | list[_PivotAggFunc] | Mapping[Hashable, _PivotAggFunc]
+        _PivotAggFunc | Sequence[_PivotAggFunc] | Mapping[Hashable, _PivotAggFunc]
     ) = ...,
     fill_value: Scalar | None = ...,
     margins: bool = ...,
@@ -111,9 +114,9 @@ def pivot_table(
 def pivot(
     data: DataFrame,
     *,
-    index: _NonIterableHashable | list[HashableT1] = ...,
-    columns: _NonIterableHashable | list[HashableT2] = ...,
-    values: _NonIterableHashable | list[HashableT3] = ...,
+    index: _NonIterableHashable | Sequence[HashableT1] = ...,
+    columns: _NonIterableHashable | Sequence[HashableT2] = ...,
+    values: _NonIterableHashable | Sequence[HashableT3] = ...,
 ) -> DataFrame: ...
 @overload
 def crosstab(

tests/test_frame.py

@@ -1337,6 +1337,17 @@ def test_types_pivot_table() -> None:
         ),
         pd.DataFrame,
     )
+    check(
+        assert_type(
+            df.pivot_table(
+                index=df["col1"].name,
+                columns=df["col3"].name,
+                values=[df["col2"].name, df["col4"].name],
+            ),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
 
 
 def test_pivot_table_sort():