Type generic pipe with function params

Author: paw-lu

pandas/_typing.py

@@ -90,16 +90,23 @@
     from typing import SupportsIndex
 
     if sys.version_info >= (3, 10):
-        from typing import TypeGuard  # pyright: ignore[reportUnusedImport]
+        from typing import (  # pyright: ignore[reportUnusedImport]
+            ParamSpec,
+            TypeGuard,
+        )
     else:
-        from typing_extensions import TypeGuard  # pyright: ignore[reportUnusedImport]
+        from typing_extensions import (  # pyright: ignore[reportUnusedImport]
+            ParamSpec,
+            TypeGuard,
+        )
 
     if sys.version_info >= (3, 11):
         from typing import Self  # pyright: ignore[reportUnusedImport]
     else:
         from typing_extensions import Self  # pyright: ignore[reportUnusedImport]
 else:
     npt: Any = None
+    ParamSpec: Any = None
     Self: Any = None
     TypeGuard: Any = None
 
@@ -222,6 +229,7 @@ def __reversed__(self) -> Iterator[_T_co]:
 
 # to maintain type information across generic functions and parametrization
 T = TypeVar("T")
+P = ParamSpec("P")
 
 # used in decorators to preserve the signature of the function it decorates
 # see https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators

pandas/core/generic.py

@@ -213,6 +213,7 @@
     )
 
     from pandas._libs.tslibs import BaseOffset
+    from pandas._typing import P
 
     from pandas import (
         DataFrame,
@@ -223,6 +224,12 @@
     from pandas.core.indexers.objects import BaseIndexer
     from pandas.core.resample import Resampler
 
+    if sys.version_info >= (3, 10):
+        from typing import Concatenate
+
+    else:
+        from typing_extensions import Concatenate
+
 # goal is to be able to define the docs close to function, while still being
 # able to share
 _shared_docs = {**_shared_docs}
@@ -6118,13 +6125,31 @@ def sample(
 
         return result
 
+    @overload
+    def pipe(
+        self,
+        func: Callable[Concatenate[Self, P], T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T:
+        ...
+
+    @overload
+    def pipe(
+        self,
+        func: tuple[Callable[..., T], str],
+        *args: Any,
+        **kwargs: Any,
+    ) -> T:
+        ...
+
     @final
     @doc(klass=_shared_doc_kwargs["klass"])
     def pipe(
         self,
-        func: Callable[..., T] | tuple[Callable[..., T], str],
-        *args,
-        **kwargs,
+        func: Callable[Concatenate[Self, P], T] | tuple[Callable[..., T], str],
+        *args: Any,
+        **kwargs: Any,
     ) -> T:
         r"""
         Apply chainable functions that expect Series or DataFrames.