Skip to content

Commit f89d534

Browse files
committed
Type generic pipe with function params
1 parent fa5e835 commit f89d534

File tree

2 files changed

+35
-4
lines changed

2 files changed

+35
-4
lines changed

pandas/_typing.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,18 +90,29 @@
9090
from typing import SupportsIndex
9191

9292
if sys.version_info >= (3, 10):
93+
from typing import Concatenate # pyright: ignore[reportUnusedImport]
94+
from typing import ParamSpec
9395
from typing import TypeGuard # pyright: ignore[reportUnusedImport]
9496
else:
95-
from typing_extensions import TypeGuard # pyright: ignore[reportUnusedImport]
97+
from typing_extensions import ( # pyright: ignore[reportUnusedImport]
98+
Concatenate,
99+
ParamSpec,
100+
TypeGuard,
101+
)
102+
103+
P = ParamSpec("P")
96104

97105
if sys.version_info >= (3, 11):
98106
from typing import Self # pyright: ignore[reportUnusedImport]
99107
else:
100108
from typing_extensions import Self # pyright: ignore[reportUnusedImport]
109+
101110
else:
102111
npt: Any = None
112+
ParamSpec: Any = None
103113
Self: Any = None
104114
TypeGuard: Any = None
115+
Concatenate: Any = None
105116

106117
HashableT = TypeVar("HashableT", bound=Hashable)
107118
MutableMappingT = TypeVar("MutableMappingT", bound=MutableMapping)

pandas/core/generic.py

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
Axis,
5151
AxisInt,
5252
CompressionOptions,
53+
Concatenate,
5354
DtypeArg,
5455
DtypeBackend,
5556
DtypeObj,
@@ -213,6 +214,7 @@
213214
)
214215

215216
from pandas._libs.tslibs import BaseOffset
217+
from pandas._typing import P
216218

217219
from pandas import (
218220
DataFrame,
@@ -6118,13 +6120,31 @@ def sample(
61186120

61196121
return result
61206122

6123+
@overload
6124+
def pipe(
6125+
self,
6126+
func: Callable[Concatenate[Self, P], T],
6127+
*args: P.args,
6128+
**kwargs: P.kwargs,
6129+
) -> T:
6130+
...
6131+
6132+
@overload
6133+
def pipe(
6134+
self,
6135+
func: tuple[Callable[..., T], str],
6136+
*args: Any,
6137+
**kwargs: Any,
6138+
) -> T:
6139+
...
6140+
61216141
@final
61226142
@doc(klass=_shared_doc_kwargs["klass"])
61236143
def pipe(
61246144
self,
6125-
func: Callable[..., T] | tuple[Callable[..., T], str],
6126-
*args,
6127-
**kwargs,
6145+
func: Callable[Concatenate[Self, P], T] | tuple[Callable[..., T], str],
6146+
*args: Any,
6147+
**kwargs: Any,
61286148
) -> T:
61296149
r"""
61306150
Apply chainable functions that expect Series or DataFrames.

0 commit comments

Comments
 (0)