TYP: GH1654 concat / pivot_table and plotting fix for 3.0 (#1720)

* GH1654 Concat/Pivot_table and plotting fix for 3.0

* GH1654 Series.map supports engine

* GH1654 Add passing a series to json_normalize

* GH1654 PR Feedback

* GH1654 PR Feedback

Author: loicdiridollou

pandas-stubs/core/frame.pyi

@@ -1364,6 +1364,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
         margins_name: _str = "All",
         observed: _bool = True,
         sort: _bool = True,
+        **kwargs: Any,
     ) -> Self: ...
     def stack(
         self,
@@ -1563,6 +1564,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
         self,
         func: Callable[..., Any],
         na_action: Literal["ignore"] | None = None,
+        engine: Any = None,
         **kwargs: Any,
     ) -> Self: ...
     def join(

pandas-stubs/core/reshape/concat.pyi

@@ -40,6 +40,19 @@ def concat(
     sort: bool = False,
 ) -> Never: ...
 @overload
+def concat(
+    objs: Iterable[NDFrame | None] | Mapping[HashableT1, NDFrame | None],
+    *,
+    axis: Axis = 0,
+    join: Literal["inner", "outer"] = "outer",
+    ignore_index: Literal[True],
+    keys: Iterable[HashableT2],
+    levels: Sequence[list[HashableT3] | tuple[HashableT3, ...]] | None = None,
+    names: list[HashableT4] | None = None,
+    verify_integrity: bool = False,
+    sort: bool = False,
+) -> Never: ...
+@overload
 def concat(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     objs: Iterable[Series[S2] | None] | Mapping[HashableT1, Series[S2] | None],
     *,

pandas-stubs/core/reshape/pivot.pyi

@@ -61,6 +61,7 @@ def pivot_table(
     margins_name: Hashable = "All",
     observed: bool = True,
     sort: bool = True,
+    **kwargs: Any,
 ) -> DataFrame: ...
 
 # Can only use Index or ndarray when index or columns is a Grouper
@@ -78,6 +79,7 @@ def pivot_table(
     margins_name: Hashable = "All",
     observed: bool = True,
     sort: bool = True,
+    **kwargs: Any,
 ) -> DataFrame: ...
 @overload
 def pivot_table(
@@ -93,6 +95,7 @@ def pivot_table(
     margins_name: Hashable = "All",
     observed: bool = True,
     sort: bool = True,
+    **kwargs: Any,
 ) -> DataFrame: ...
 def pivot(
     data: DataFrame,

pandas-stubs/core/series.pyi

@@ -1116,39 +1116,45 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
         self,
         func: Callable[Concatenate[S1, ...], S2 | NAType],
         na_action: Literal["ignore"],
+        engine: Callable[..., Any] | None = None,
         **kwargs: Any,
     ) -> Series[S2]: ...
     @overload
     def map(
         self,
         func: Mapping[S1, S2] | Series[S2],
         na_action: Literal["ignore"],
+        engine: None = None,
     ) -> Series[S2]: ...
     @overload
     def map(
         self,
         func: Callable[Concatenate[S1 | NAType, ...], S2 | NAType],
         na_action: None = None,
+        engine: Callable[..., Any] | None = None,
         **kwargs: Any,
     ) -> Series[S2]: ...
     @overload
     def map(
         self,
         func: Mapping[S1, S2] | Series[S2],
         na_action: None = None,
+        engine: None = None,
     ) -> Series[S2]: ...
     @overload
     def map(
         self,
         func: Callable[..., Any],
         na_action: Literal["ignore"] | None = None,
+        engine: Callable[..., Any] | None = None,
         **kwargs: Any,
     ) -> Series: ...
     @overload
     def map(
         self,
         func: Mapping[Any, Any] | Series,
         na_action: Literal["ignore"] | None = None,
+        engine: None = None,
     ) -> Series: ...
     @overload
     def aggregate(

pandas-stubs/core/window/ewm.pyi

@@ -54,13 +54,13 @@ class ExponentialMovingWindow(BaseWindow[NDFrameT]):
         **kwargs: Any,
     ) -> Series: ...
     @overload
-    def aggregate(  # ty: ignore[invalid-method-override]  # pyright: ignore[reportIncompatibleMethodOverride]
+    def aggregate(  # pyright: ignore[reportIncompatibleMethodOverride]
         self: BaseWindow[DataFrame],
         func: str,
         *args: Any,
         **kwargs: Any,
     ) -> DataFrame: ...
-    agg = aggregate  # type: ignore[assignment]  # ty: ignore[invalid-method-override]  # pyrefly: ignore[bad-override]
+    agg = aggregate  # type: ignore[assignment] # pyrefly: ignore[bad-override]
 
 class ExponentialMovingWindowGroupby(
     BaseWindowGroupby[NDFrameT], ExponentialMovingWindow[NDFrameT]

pandas-stubs/io/json/_normalize.pyi

@@ -1,11 +1,13 @@
+from collections.abc import Iterable
 from typing import Any
 
 from pandas import DataFrame
+from pandas.core.series import Series
 
 from pandas._typing import IgnoreRaise
 
 def json_normalize(
-    data: dict[str, Any] | list[dict[str, Any]],
+    data: dict[str, Any] | Iterable[dict[str, Any]] | Series,
     record_path: str | list[str] | None = None,
     meta: str | list[str | list[str]] | None = None,
     meta_prefix: str | None = None,

pandas-stubs/plotting/_core.pyi

@@ -468,6 +468,7 @@ class PlotAccessor:
             | Callable[[gaussian_kde], float]
             | None
         ) = ...,
+        weights: np_ndarray_float | Series[float] | None = None,
         ind: np_ndarray_float | int | None = ...,
         *,
         subplots: Literal[False] | None = ...,
@@ -482,6 +483,7 @@ class PlotAccessor:
             | Callable[[gaussian_kde], float]
             | None
         ) = ...,
+        weights: np_ndarray_float | Series[float] | None = None,
         ind: np_ndarray_float | int | None = ...,
         *,
         subplots: Literal[True],

tests/frame/test_frame.py

@@ -1527,6 +1527,25 @@ def test_types_pivot_table() -> None:
         pd.DataFrame,
     )
 
+    # test passing kwargs for the aggfunc to pivot_table
+    df = pd.DataFrame(
+        {
+            "A": ["good", "bad", "good", "bad", "good"],
+            "B": ["one", "two", "one", "three", "two"],
+            "X": [2, 5, 4, 20, 10],
+        }
+    )
+
+    check(
+        assert_type(
+            pd.pivot_table(
+                df, index="A", columns="B", values="X", aggfunc="std", ddof=2
+            ),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
+
 
 def test_pivot_table_aggfunc_string_reduction(sample_df: pd.DataFrame) -> None:
     """Test string aggfunc with reduction functions from ReductionKernelType."""

tests/series/test_series.py

@@ -2783,6 +2783,50 @@ def func(x: int, y: int) -> int:
     )
 
 
+def test_map_engine() -> None:
+    s: pd.Series[int] = pd.Series([1, 2, 3])
+
+    def to_str(x: int) -> str:
+        return str(x)
+
+    # engine=None with na_action="ignore": func takes S1, returns S2
+    check(
+        assert_type(s.map(to_str, na_action="ignore", engine=None), "pd.Series[str]"),
+        pd.Series,
+        str,
+    )
+
+    # engine=None with na_action=None: func must handle NAType
+    def nullable_to_str(x: int | NAType) -> str | NAType:
+        return str(x) if isinstance(x, int) else x
+
+    check(
+        assert_type(
+            s.map(nullable_to_str, na_action=None, engine=None), "pd.Series[str]"
+        ),
+        pd.Series,
+        str,
+    )
+
+    # Mapping/Series: engine must be None (omitted or explicit)
+    mapping = {1: "a", 2: "b", 3: "c"}
+    check(
+        assert_type(s.map(mapping, engine=None), "pd.Series[str]"),
+        pd.Series,
+        str,
+    )
+
+    if TYPE_CHECKING:
+        # engine accepts any callable (e.g. a JIT decorator like numba.jit).
+        # Cannot be tested at runtime: pandas validates __pandas_udf__ on the engine.
+        def jit_decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+            return f
+
+        assert_type(
+            s.map(to_str, na_action="ignore", engine=jit_decorator), "pd.Series[str]"
+        )
+
+
 def test_case_when() -> None:
     c = pd.Series([6, 7, 8, 9], name="c")
     a = pd.Series([0, 0, 1, 2])

tests/test_pandas.py

@@ -403,6 +403,15 @@ def test_concat_args() -> None:
         pd.DataFrame,
     )
 
+    if TYPE_CHECKING_INVALID_USAGE:
+
+        def _cannot_pass_ignore_index_eq_true_and_keys() -> (  # pyright: ignore[reportUnusedFunction]
+            None
+        ):
+            assert_type(
+                pd.concat([df, df2], ignore_index=True, keys=["df1", "df2"]), Never
+            )
+
 
 def test_frame_subclass_concat() -> None:
     """Test concatenate subclass of DataFrame GH1396."""
@@ -445,13 +454,37 @@ def test_types_json_normalize() -> None:
         ),
         pd.DataFrame,
     )
+
+    # iterable case
+    check(
+        assert_type(pd.json_normalize(data=(dic for dic in data1)), pd.DataFrame),
+        pd.DataFrame,
+    )
+
     data2: dict[str, Any] = {"name": {"given": "Mose", "family": "Regner"}}
     check(
         assert_type(data2, dict[str, Any]),
         dict,
     )
     check(assert_type(pd.json_normalize(data=data2), pd.DataFrame), pd.DataFrame)
 
+    # series case
+    data = [
+        {
+            "id": 1,
+            "name": "Cole Volk",
+            "fitness": {"height": 130, "weight": 60},
+        },
+        {"name": "Mark Reg", "fitness": {"height": 130, "weight": 60}},
+        {
+            "id": 2,
+            "name": "Faye Raker",
+            "fitness": {"height": 130, "weight": 60},
+        },
+    ]
+    series = pd.Series(data, index=pd.Index(["a", "b", "c"]))
+    check(assert_type(pd.json_normalize(series), pd.DataFrame), pd.DataFrame)
+
 
 def test_isna() -> None:
     # https://github.com/pandas-dev/pandas-stubs/issues/264
@@ -2128,6 +2161,24 @@ def g(x: pd.Series) -> int:
         index=idx,
     )
 
+    df = pd.DataFrame(
+        {
+            "A": ["good", "bad", "good", "bad", "good"],
+            "B": ["one", "two", "one", "three", "two"],
+            "X": [2, 5, 4, 20, 10],
+        }
+    )
+
+    check(
+        assert_type(
+            pd.pivot_table(
+                df, index="A", columns="B", values="X", aggfunc="std", ddof=2
+            ),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
+
 
 def test_pivot_table_aggfunc_string_reduction(sample_df: pd.DataFrame) -> None:
     """Test string aggfunc with reduction functions from ReductionKernelType."""