Type generic pipe with function params

Author: paw-lu

pandas/_typing.py

@@ -90,9 +90,15 @@
  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]
@@ -222,6 +228,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.