ENH: Type pipe with ParamSpec

[paw-lu]

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.

pandas/pandas/core/generic.py

Lines 6039 to 6044 in 2a953cf

	def pipe(
	self,
	func: Callable[..., T] | tuple[Callable[..., T], str],
	*args,
	**kwargs,
	) -> T:

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.

Additional Context

pipe in pandas-stubs.

[twoertwein]

I think it would be best first to have a pandas-stubs PR (with typing tests) and after that is merged and works well a pandas PR.

@Dr-Irv

[Dr-Irv]

I think it would be best first to have a pandas-stubs PR (with typing tests) and after that is merged and works well a pandas PR.

@Dr-Irv

I agree. The goal is to improve type checking for users, which is why we have pandas-stubs .

Also, see pandas-dev/pandas-stubs#738

[paw-lu]

Alright, sounds good. I'll continue this in pandas-dev/pandas-stubs, then follow up here if things go well there. Thanks!

[paw-lu]

This was implemented and merged in

• Type args and kwargs in pipe pandas-stubs#823

Would like to port this over into pandas-dev, but I think the minimum python version here is 3.9. And I don't think typing_extensions is an existing dependency. This means that we can't use ParamSpec here right (which was introduced in 3.10)? Which means this can't be implemented here until it moves to 3.10.

EDIT: My mistake, was on a old local version, I can see now typing_extensions was added as a dev dependency. This can work now, I'll get started on a draft PR using typing_extensions.