Skip to content

Commit 528dc1e

Browse files
committed
Type groupby pipe with function params
1 parent f7e3464 commit 528dc1e

File tree

1 file changed

+29
-5
lines changed

1 file changed

+29
-5
lines changed

pandas/core/groupby/groupby.py

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ class providing the base-class of operations.
2929
Union,
3030
cast,
3131
final,
32+
overload,
3233
)
3334
import warnings
3435

@@ -55,7 +56,6 @@ class providing the base-class of operations.
5556
PositionalIndexer,
5657
RandomState,
5758
Scalar,
58-
T,
5959
npt,
6060
)
6161
from pandas.compat.numpy import function as nv
@@ -147,7 +147,13 @@ class providing the base-class of operations.
147147
)
148148

149149
if TYPE_CHECKING:
150-
from typing import Any
150+
from pandas._typing import (
151+
Any,
152+
Concatenate,
153+
P,
154+
Self,
155+
T,
156+
)
151157

152158
from pandas.core.resample import Resampler
153159
from pandas.core.window import (
@@ -988,6 +994,24 @@ def _selected_obj(self):
988994
def _dir_additions(self) -> set[str]:
989995
return self.obj._dir_additions()
990996

997+
@overload
998+
def pipe(
999+
self,
1000+
func: Callable[Concatenate[Self, P], T],
1001+
*args: P.args,
1002+
**kwargs: P.kwargs,
1003+
) -> T:
1004+
...
1005+
1006+
@overload
1007+
def pipe(
1008+
self,
1009+
func: tuple[Callable[..., T], str],
1010+
*args: Any,
1011+
**kwargs: Any,
1012+
) -> T:
1013+
...
1014+
9911015
@Substitution(
9921016
klass="GroupBy",
9931017
examples=dedent(
@@ -1013,9 +1037,9 @@ def _dir_additions(self) -> set[str]:
10131037
@Appender(_pipe_template)
10141038
def pipe(
10151039
self,
1016-
func: Callable[..., T] | tuple[Callable[..., T], str],
1017-
*args,
1018-
**kwargs,
1040+
func: Callable[Concatenate[Self, P], T] | tuple[Callable[..., T], str],
1041+
*args: Any,
1042+
**kwargs: Any,
10191043
) -> T:
10201044
return com.pipe(self, func, *args, **kwargs)
10211045

0 commit comments

Comments
 (0)