Skip to content

Commit 519d0e7

Browse files
committed
Type common pipe with function params
1 parent 6155d05 commit 519d0e7

File tree

2 files changed

+31
-3
lines changed

2 files changed

+31
-3
lines changed

pandas/_typing.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,7 @@ def __reversed__(self) -> Iterator[_T_co]:
233233

234234
# to maintain type information across generic functions and parametrization
235235
T = TypeVar("T")
236+
S = TypeVar("S")
236237

237238
# used in decorators to preserve the signature of the function it decorates
238239
# see https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators

pandas/core/common.py

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,11 @@
5151
from pandas._typing import (
5252
AnyArrayLike,
5353
ArrayLike,
54+
Concatenate,
5455
NpDtype,
56+
P,
5557
RandomState,
58+
S,
5659
T,
5760
)
5861

@@ -463,8 +466,31 @@ def random_state(state: RandomState | None = None):
463466
)
464467

465468

469+
@overload
470+
def pipe(
471+
obj: S,
472+
func: Callable[Concatenate[S, P], T],
473+
*args: P.args,
474+
**kwargs: P.kwargs,
475+
) -> T:
476+
...
477+
478+
479+
@overload
480+
def pipe(
481+
obj: Any,
482+
func: tuple[Callable[..., T], str],
483+
*args: Any,
484+
**kwargs: Any,
485+
) -> T:
486+
...
487+
488+
466489
def pipe(
467-
obj, func: Callable[..., T] | tuple[Callable[..., T], str], *args, **kwargs
490+
obj: S,
491+
func: Callable[Concatenate[S, P], T] | tuple[Callable[..., T], str],
492+
*args: Any,
493+
**kwargs: Any,
468494
) -> T:
469495
"""
470496
Apply a function ``func`` to object ``obj`` either by passing obj as the
@@ -490,12 +516,13 @@ def pipe(
490516
object : the return type of ``func``.
491517
"""
492518
if isinstance(func, tuple):
493-
func, target = func
519+
# Assigning to func_ so pyright understands that it's a callable
520+
func_, target = func
494521
if target in kwargs:
495522
msg = f"{target} is both the pipe target and a keyword argument"
496523
raise ValueError(msg)
497524
kwargs[target] = obj
498-
return func(*args, **kwargs)
525+
return func_(*args, **kwargs)
499526
else:
500527
return func(obj, *args, **kwargs)
501528

0 commit comments

Comments
 (0)