From f89d534c95be05ef975daa874d0963ecc564b4d3 Mon Sep 17 00:00:00 2001 From: "Paulo S. Costa" Date: Thu, 4 Jan 2024 22:37:20 -0800 Subject: [PATCH 1/5] Type generic pipe with function params --- pandas/_typing.py | 13 ++++++++++++- pandas/core/generic.py | 26 +++++++++++++++++++++++--- 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/pandas/_typing.py b/pandas/_typing.py index a80f9603493a7..fa9dc14bb4bd7 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -90,18 +90,29 @@ from typing import SupportsIndex if sys.version_info >= (3, 10): + from typing import Concatenate # pyright: ignore[reportUnusedImport] + from typing import ParamSpec from typing import TypeGuard # pyright: ignore[reportUnusedImport] else: - from typing_extensions import TypeGuard # pyright: ignore[reportUnusedImport] + from typing_extensions import ( # pyright: ignore[reportUnusedImport] + Concatenate, + ParamSpec, + TypeGuard, + ) + + P = ParamSpec("P") 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 + Concatenate: Any = None HashableT = TypeVar("HashableT", bound=Hashable) MutableMappingT = TypeVar("MutableMappingT", bound=MutableMapping) diff --git a/pandas/core/generic.py b/pandas/core/generic.py index b37f22339fcfd..40ce9499d69c2 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -50,6 +50,7 @@ Axis, AxisInt, CompressionOptions, + Concatenate, DtypeArg, DtypeBackend, DtypeObj, @@ -213,6 +214,7 @@ ) from pandas._libs.tslibs import BaseOffset + from pandas._typing import P from pandas import ( DataFrame, @@ -6118,13 +6120,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. From d4833701a2b73d148c074c9d91430b8c141038d7 Mon Sep 17 00:00:00 2001 From: "Paulo S. Costa" Date: Sun, 7 Jan 2024 12:23:34 -0800 Subject: [PATCH 2/5] Type common pipe with function params --- pandas/core/common.py | 36 +++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/pandas/core/common.py b/pandas/core/common.py index 7d864e02be54e..69b602feee3ea 100644 --- a/pandas/core/common.py +++ b/pandas/core/common.py @@ -24,6 +24,7 @@ TYPE_CHECKING, Any, Callable, + TypeVar, cast, overload, ) @@ -51,7 +52,9 @@ from pandas._typing import ( AnyArrayLike, ArrayLike, + Concatenate, NpDtype, + P, RandomState, T, ) @@ -463,8 +466,34 @@ def random_state(state: RandomState | None = None): ) +_T = TypeVar("_T") # Secondary TypeVar for use in pipe's type hints + + +@overload +def pipe( + obj: _T, + func: Callable[Concatenate[_T, 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: _T, + func: Callable[Concatenate[_T, P], T] | tuple[Callable[..., T], str], + *args: Any, + **kwargs: Any, ) -> T: """ Apply a function ``func`` to object ``obj`` either by passing obj as the @@ -490,12 +519,13 @@ def pipe( object : the return type of ``func``. """ if isinstance(func, tuple): - func, target = func + # Assigning to func_ so pyright understands that it's a callable + func_, target = func if target in kwargs: msg = f"{target} is both the pipe target and a keyword argument" raise ValueError(msg) kwargs[target] = obj - return func(*args, **kwargs) + return func_(*args, **kwargs) else: return func(obj, *args, **kwargs) From 75889e3a8171ea35f4f358223b083262315ec1c7 Mon Sep 17 00:00:00 2001 From: "Paulo S. Costa" Date: Sun, 7 Jan 2024 12:30:57 -0800 Subject: [PATCH 3/5] Type resample pipe with function params --- pandas/core/resample.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/pandas/core/resample.py b/pandas/core/resample.py index 31309777c154d..924f9e6d49040 100644 --- a/pandas/core/resample.py +++ b/pandas/core/resample.py @@ -9,6 +9,7 @@ cast, final, no_type_check, + overload, ) import warnings @@ -97,12 +98,16 @@ from collections.abc import Hashable from pandas._typing import ( + Any, AnyArrayLike, Axis, AxisInt, + Concatenate, Frequency, IndexLabel, InterpolateOptions, + P, + Self, T, TimedeltaConvertibleTypes, TimeGrouperOrigin, @@ -254,6 +259,24 @@ def _get_binner(self): bin_grouper = BinGrouper(bins, binlabels, indexer=self._indexer) return binner, bin_grouper + @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 @Substitution( klass="Resampler", @@ -278,9 +301,9 @@ def _get_binner(self): @Appender(_pipe_template) 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: return super().pipe(func, *args, **kwargs) From d413a69c8ab97abab2273fe2110043d4147074c3 Mon Sep 17 00:00:00 2001 From: "Paulo S. Costa" Date: Sun, 7 Jan 2024 12:41:45 -0800 Subject: [PATCH 4/5] Type groupby pipe with function params --- pandas/core/groupby/groupby.py | 34 +++++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index c9beaee55d608..f1ca05a312d2a 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -29,6 +29,7 @@ class providing the base-class of operations. Union, cast, final, + overload, ) import warnings @@ -55,7 +56,6 @@ class providing the base-class of operations. PositionalIndexer, RandomState, Scalar, - T, npt, ) from pandas.compat.numpy import function as nv @@ -147,7 +147,13 @@ class providing the base-class of operations. ) if TYPE_CHECKING: - from typing import Any + from pandas._typing import ( + Any, + Concatenate, + P, + Self, + T, + ) from pandas.core.resample import Resampler from pandas.core.window import ( @@ -988,6 +994,24 @@ def _selected_obj(self): def _dir_additions(self) -> set[str]: return self.obj._dir_additions() + @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: + ... + @Substitution( klass="GroupBy", examples=dedent( @@ -1013,9 +1037,9 @@ def _dir_additions(self) -> set[str]: @Appender(_pipe_template) 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: return com.pipe(self, func, *args, **kwargs) From b1a3fc406e036bc3a6ef3d6eff5bd649b4bede93 Mon Sep 17 00:00:00 2001 From: "Paulo S. Costa" Date: Sun, 7 Jan 2024 12:56:45 -0800 Subject: [PATCH 5/5] Type style pipe function params and tuple func --- pandas/io/formats/style.py | 31 +++++++++++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/pandas/io/formats/style.py b/pandas/io/formats/style.py index b62f7581ac220..db4a7e42f3bd3 100644 --- a/pandas/io/formats/style.py +++ b/pandas/io/formats/style.py @@ -9,7 +9,6 @@ import operator from typing import ( TYPE_CHECKING, - Any, Callable, overload, ) @@ -66,15 +65,20 @@ from matplotlib.colors import Colormap from pandas._typing import ( + Any, Axis, AxisInt, + Concatenate, FilePath, IndexLabel, IntervalClosedType, Level, + P, QuantileInterpolation, Scalar, + Self, StorageOptions, + T, WriteBuffer, WriteExcelBuffer, ) @@ -3614,7 +3618,30 @@ class MyStyler(cls): # type: ignore[valid-type,misc] return MyStyler - def pipe(self, func: Callable, *args, **kwargs): + @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: + ... + + def pipe( + self, + func: Callable[Concatenate[Self, P], T] | tuple[Callable[..., T], str], + *args: Any, + **kwargs: Any, + ) -> T: """ Apply ``func(self, *args, **kwargs)``, and return the result.