ENH: Type pipe with ParamSpec #56359
Description
Feature Type
-
Adding new functionality to pandas
-
Changing existing functionality in pandas
-
Removing existing functionality in pandas
Problem Description
Right now the *args and **kwargs in pandas.DataFrame.pipe are not typed, removing any type checking benefit on the parameters if a typed function is used via pipe.
Lines 6039 to 6044 in 2a953cf
Feature Description
I'd be happy to contribute this if this seems acceptable.
I have a basic implementation here of what I'd want to apply to the real method. I'm guessing this would go in pandas-stubs as well?
One con here is that the tuple version of the func parameter does not use ParamSpec, since Concatenate requires that the additional type comes first. So users using a tuple for func wouldn't benefit from the type checking. Unfortunately, don't really know a way around this right now.
from __future__ import annotations
import dataclasses
import typing
from typing import Callable, Concatenate
P = typing.ParamSpec("P")
T = typing.TypeVar("T")
@dataclasses.dataclass
class BabyDataFrame:
data: int
def pipe(
self,
func: Callable[Concatenate[BabyDataFrame, P], T] | tuple[Callable[..., T], str],
*args: P.args,
**kwargs: P.kwargs,
) -> T:
if isinstance(func, tuple):
return func[0](*args, **kwargs, **{func[1]: self})
else:
return func(self, *args, **kwargs)
def foo(baby_data_frame: BabyDataFrame, x: int) -> int:
return baby_data_frame.data + x
def bar(x: int, baby_data_frame: BabyDataFrame) -> int:
return baby_data_frame.data + x
print(BabyDataFrame(2).pipe(foo, 1))
print(BabyDataFrame(2).pipe((bar, "baby_data_frame"), 2))Alternative Solutions
Not sure of an alternative solution here besides just implementing in pandas-stubs. Figured I might as well do both.