Type common pipe with function params

Author: paw-lu

pandas/_typing.py

@@ -233,6 +233,7 @@ def __reversed__(self) -> Iterator[_T_co]:
 
 # to maintain type information across generic functions and parametrization
 T = TypeVar("T")
+S = TypeVar("S")
 
 # 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/common.py

@@ -51,8 +51,11 @@
     from pandas._typing import (
         AnyArrayLike,
         ArrayLike,
+        Concatenate,
         NpDtype,
+        P,
         RandomState,
+        S,
         T,
     )
 
@@ -463,8 +466,31 @@ def random_state(state: RandomState | None = None):
         )
 
 
+@overload
+def pipe(
+    obj: S,
+    func: Callable[Concatenate[S, P], T],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> T:
+    ...
+
+
+@overload
+def pipe(
+    obj: Any,
+    func: tuple[Callable[..., T], str],
+    *args: Any,
+    **kwargs: Any,
+) -> T:
+    ...
+
+
 def pipe(
-    obj, func: Callable[..., T] | tuple[Callable[..., T], str], *args, **kwargs
+    obj: S,
+    func: Callable[Concatenate[S, P], T] | tuple[Callable[..., T], str],
+    *args: Any,
+    **kwargs: Any,
 ) -> T:
     """
     Apply a function ``func`` to object ``obj`` either by passing obj as the