Skip to content

Commit

Permalink
Type groupby pipe with function params
Browse files Browse the repository at this point in the history
  • Loading branch information
paw-lu committed Jan 10, 2024
1 parent 5e52626 commit 4609b8c
Showing 1 changed file with 29 additions and 5 deletions.
34 changes: 29 additions & 5 deletions pandas/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ class providing the base-class of operations.
Union,
cast,
final,
overload,
)
import warnings

Expand All @@ -55,7 +56,6 @@ class providing the base-class of operations.
PositionalIndexer,
RandomState,
Scalar,
T,
npt,
)
from pandas.compat.numpy import function as nv
Expand Down Expand Up @@ -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 (
Expand Down Expand Up @@ -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(
Expand All @@ -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)

Expand Down

0 comments on commit 4609b8c

Please sign in to comment.