Skip to content

Commit

Permalink
Type resample pipe with function params
Browse files Browse the repository at this point in the history
  • Loading branch information
paw-lu committed Jan 11, 2024
1 parent 408214d commit f7744ae
Showing 1 changed file with 26 additions and 3 deletions.
29 changes: 26 additions & 3 deletions pandas/core/resample.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
cast,
final,
no_type_check,
overload,
)
import warnings

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand All @@ -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)

Expand Down

0 comments on commit f7744ae

Please sign in to comment.