Skip to content

Commit

Permalink
Type generic pipe with function params
Browse files Browse the repository at this point in the history
  • Loading branch information
paw-lu committed Jan 7, 2024
1 parent b2bca5e commit fc6d249
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 5 deletions.
12 changes: 10 additions & 2 deletions pandas/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
31 changes: 28 additions & 3 deletions pandas/core/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,7 @@
)

from pandas._libs.tslibs import BaseOffset
from pandas._typing import P

from pandas import (
DataFrame,
Expand All @@ -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}
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit fc6d249

Please sign in to comment.