TYP: GH1654 pipe method (#1665)

* GH1654 Pipe method

* GH1654 PR Feedback

* GH1654 PR Feedback

Author: loicdiridollou

pandas-stubs/api/typing/__init__.pyi

@@ -2,6 +2,7 @@ from pandas.core.groupby import (
     DataFrameGroupBy as DataFrameGroupBy,
     SeriesGroupBy as SeriesGroupBy,
 )
+from pandas.core.indexes.frozen import FrozenList as FrozenList
 from pandas.core.resample import (
     DatetimeIndexResamplerGroupby as DatetimeIndexResamplerGroupby,
     PeriodIndexResamplerGroupby as PeriodIndexResamplerGroupby,

pandas-stubs/core/accessor.pyi

@@ -7,3 +7,5 @@ class PandasDelegate: ...
 def register_dataframe_accessor(name: str) -> Callable[[TypeT], TypeT]: ...
 def register_series_accessor(name: str) -> Callable[[TypeT], TypeT]: ...
 def register_index_accessor(name: str) -> Callable[[TypeT], TypeT]: ...
+
+class DirNamesMixin: ...

pandas-stubs/core/base.pyi

@@ -18,6 +18,7 @@ from typing import (
 import numpy as np
 import numpy.typing as npt
 from pandas._stubs_only import T_contra
+from pandas.core.accessor import DirNamesMixin
 from pandas.core.arraylike import OpsMixin
 from pandas.core.arrays import ExtensionArray
 from pandas.core.arrays.categorical import Categorical
@@ -57,6 +58,8 @@ T_INTERVAL_NP = TypeVar("T_INTERVAL_NP", bound=np.bytes_ | np.str_)
 class NoNewAttributesMixin:
     def __setattr__(self, key: str, value: Any) -> None: ...
 
+class PandasObject(DirNamesMixin): ...
+
 class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
     __array_priority__: int = ...
     @property

pandas-stubs/core/groupby/base.pyi

@@ -32,6 +32,7 @@ ReductionKernelType: TypeAlias = Literal[
     "sem",
     "size",
     "skew",
+    "kurt",
     "std",
     "sum",
     "var",

pandas-stubs/core/indexes/base.pyi

@@ -58,6 +58,7 @@ from pandas.core.base import (
 from pandas.core.frame import DataFrame
 from pandas.core.indexes.category import CategoricalIndex
 from pandas.core.indexes.datetimes import DatetimeIndex
+from pandas.core.indexes.frozen import FrozenList
 from pandas.core.indexes.interval import IntervalIndex
 from pandas.core.indexes.multi import MultiIndex
 from pandas.core.indexes.period import PeriodIndex
@@ -413,7 +414,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     @name.setter
     def name(self, value: Hashable) -> None: ...
     @property
-    def names(self) -> list[Hashable | None]: ...
+    def names(self) -> FrozenList[Hashable | None]: ...
     @names.setter
     def names(self, names: SequenceNotStr[Hashable | None]) -> None: ...
     def set_names(

pandas-stubs/core/indexes/frozen.pyi

@@ -0,0 +1,10 @@
+from typing import (
+    Generic,
+    TypeVar,
+)
+
+from pandas.core.base import PandasObject
+
+_T = TypeVar("_T")
+
+class FrozenList(PandasObject, list[_T], Generic[_T]): ...

pandas-stubs/core/indexes/multi.pyi

@@ -14,6 +14,7 @@ from typing import (
 
 import numpy as np
 import pandas as pd
+from pandas.api.typing import FrozenList
 from pandas.core.indexes.base import Index
 
 from pandas._typing import (
@@ -87,7 +88,7 @@ class MultiIndex(Index):
         names: SequenceNotStr[Hashable] | None = None,
     ) -> Self: ...
     @property  # Should be read-only
-    def levels(self) -> list[Index]: ...
+    def levels(self) -> FrozenList[Index]: ...
     @overload
     def set_levels(
         self,
@@ -105,7 +106,7 @@ class MultiIndex(Index):
         verify_integrity: bool = True,
     ) -> MultiIndex: ...
     @property
-    def codes(self) -> list[np_1darray_int8]: ...
+    def codes(self) -> FrozenList[np_1darray_int8]: ...
     @overload
     def set_codes(
         self,

pandas-stubs/core/window/rolling.pyi

@@ -7,6 +7,7 @@ from collections.abc import (
 import datetime as dt
 from typing import (
     Any,
+    Concatenate,
     Generic,
     Self,
     overload,
@@ -28,7 +29,9 @@ from pandas._typing import (
     CalculationMethod,
     IntervalClosedType,
     NDFrameT,
+    P,
     QuantileInterpolation,
+    T,
     WindowingEngine,
     WindowingEngineKwargs,
     WindowingRankType,
@@ -171,6 +174,20 @@ class RollingAndExpandingMixin(BaseWindow[NDFrameT]):
     def first(self, numeric_only: bool = False) -> NDFrameT: ...
     def last(self, numeric_only: bool = False) -> NDFrameT: ...
     def nunique(self, numeric_only: bool = False) -> NDFrameT: ...
+    @overload
+    def pipe(
+        self,
+        func: Callable[Concatenate[Self, P], T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T: ...
+    @overload
+    def pipe(
+        self,
+        func: tuple[Callable[Concatenate[Self, P], T], str],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T: ...
 
 class Rolling(RollingAndExpandingMixin[NDFrameT]): ...
 class RollingGroupby(BaseWindowGroupby[NDFrameT], Rolling[NDFrameT]): ...

tests/frame/test_frame.py

@@ -41,6 +41,8 @@
     DatetimeIndexResampler,
     Resampler,
 )
+from pandas.core.window.expanding import Expanding
+from pandas.core.window.rolling import Rolling
 import pytest
 import xarray as xr
 
@@ -4324,3 +4326,31 @@ def test_expanding_nunique() -> None:
         assert_type(df.expanding().nunique(numeric_only=False), pd.DataFrame),
         pd.DataFrame,
     )
+
+
+def test_frame_pipe() -> None:
+    """Test Rolling.pipe and Expanding.pipe."""
+    df = pd.DataFrame(
+        {
+            "B": np.random.default_rng(2).standard_normal(10),
+            "C": np.random.default_rng(2).standard_normal(10),
+        }
+    )
+
+    def func_r(x: Rolling[pd.DataFrame], k: int) -> pd.DataFrame:
+        return x.max() - k * x.min()
+
+    def func_e(x: Expanding[pd.DataFrame], k: int) -> pd.DataFrame:
+        return x.max() - k * x.min()
+
+    check(
+        assert_type(df.rolling(2).pipe(lambda x: x.min() - x.max()), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(assert_type(df.rolling(2).pipe(func_r, k=2), pd.DataFrame), pd.DataFrame)
+
+    check(
+        assert_type(df.expanding().pipe(lambda x: x.min() - x.max()), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(assert_type(df.expanding().pipe(func_e, k=2), pd.DataFrame), pd.DataFrame)

tests/indexes/test_indexes.py

@@ -23,6 +23,7 @@
 from pandas.core.indexes.base import Index
 from pandas.core.indexes.category import CategoricalIndex
 from pandas.core.indexes.datetimes import DatetimeIndex
+from pandas.core.indexes.frozen import FrozenList
 import pytest
 
 from tests import (
@@ -1399,20 +1400,20 @@ def test_index_naming() -> None:
     """
     Test index names type both for the getter and the setter.
     The names of an index should be settable with a sequence (not str) and names
-    property is a list[Hashable | None] (FrozenList).
+    property is a FrozenList[Hashable | None].
     """
     df = pd.DataFrame({"a": ["a", "b", "c"], "i": [10, 11, 12]})
 
     df.index.names = ["idx"]
-    check(assert_type(df.index.names, list[Hashable | None]), list)
+    check(assert_type(df.index.names, FrozenList[Hashable | None]), FrozenList)
     df.index.names = [3]
-    check(assert_type(df.index.names, list[Hashable | None]), list)
+    check(assert_type(df.index.names, FrozenList[Hashable | None]), FrozenList)
     df.index.names = ("idx2",)
-    check(assert_type(df.index.names, list[Hashable | None]), list)
+    check(assert_type(df.index.names, FrozenList[Hashable | None]), FrozenList)
     df.index.names = [None]
-    check(assert_type(df.index.names, list[Hashable | None]), list)
+    check(assert_type(df.index.names, FrozenList[Hashable | None]), FrozenList)
     df.index.names = (None,)
-    check(assert_type(df.index.names, list[Hashable | None]), list)
+    check(assert_type(df.index.names, FrozenList[Hashable | None]), FrozenList)
 
 
 def test_index_searchsorted() -> None: