Add support for arrays in lower and upper args of DataFrame.clip

[ldouteau]

• Closes lower and upper from DataFrame.clip don't accept array-like #980
• Tests added: Please use assert_type() to assert the type of any return value

[Dr-Irv]

It's possible (but not required) that you could also handle this via typing checks by creating overloads that are the acceptable combinations of the arguments lower, upper and axis

[Dr-Irv]

It's possible (but not required) that you could also handle this via typing checks by creating overloads that are the acceptable combinations of the arguments lower, upper and axis

I'm OK to approve this as is, but let me know if you want to try to make these changes I suggested as part of this PR.

[ldouteau]

I'm not sure how to type the fact that "either one of lower or upper is ArrayLike". If you have any suggestion to propose, i can add overloads

Regarding the test i think it should stay as-is, as it tests the pandas version of the function

[ldouteau]

Just thinking of this implementation. One question though: the case where upper and lower are both arrays is covered by both overloads. Is that an issue ? I'm not that familiar with overloads in general

    @overload
    def clip(
        self,
        lower: AnyArrayLike = ...,
        upper: float | AnyArrayLike | None = ...,
        *,
        axis: Axis = ...,
        inplace: _bool = ...,
        **kwargs,
    ) -> DataFrame: ...
    @overload
    def clip(
        self,
        lower: float | AnyArrayLike | None = ...,
        upper: AnyArrayLike = ...,
        *,
        axis: Axis = ...,
        inplace: _bool = ...,
        **kwargs,
    ) -> DataFrame: ...

[Dr-Irv]

I'm not sure how to type the fact that "either one of lower or upper is ArrayLike". If you have any suggestion to propose, i can add overloads

I don't think that is required. df.clip() works fine, as does df.clip(3, 4). So maybe your test isn't right?

Regarding the test i think it should stay as-is, as it tests the pandas version of the function

If you can get the overloads to work, then you'd add tests using TYPE_CHECKING_INVALID_USAGE to make sure the type checkers caught the invalid usage. But I'm now not sure the overloads are needed.

[ldouteau]

maybe your test isn't right?

Well the test is not exhaustive: df.clip([1, 2]) will pass but df.clip(pd.Series([1, 2])) will fail with a value error.

What's your opinion on this ? Getting dedicated overloads using pd.Series ?

[Dr-Irv]

Just thinking of this implementation. One question though: the case where upper and lower are both arrays is covered by both overloads. Is that an issue ?

I think the solution here is to disallow None as follows:

    @overload
    def clip(
        self,
        lower:  AnyArrayLike,
        upper: float | AnyArrayLike | None,
        axis: None,
        *,
        inplace: _bool = ...,
        **kwargs,
    ) -> Never: ...
    @overload
    def clip(
        self,
        lower:  float | AnyArrayLike | None,
        upper:  AnyArrayLike,
        axis: None,
        *,
        inplace: _bool = ...,
        **kwargs,
    ) -> Never: ...

Testing this means you have to move away from using pytest.mark.parameterize, and explicitly test each possible condition. This is preferred anyway from how we want to test the types.

You can use assert_never() to test the cases where it should fail.

[Dr-Irv]

maybe your test isn't right?

Well the test is not exhaustive: df.clip([1, 2]) will pass but df.clip(pd.Series([1, 2])) will fail with a value error.

What's your opinion on this ? Getting dedicated overloads using pd.Series ?

maybe your test isn't right?

Well the test is not exhaustive: df.clip([1, 2]) will pass but df.clip(pd.Series([1, 2])) will fail with a value error.

What's your opinion on this ? Getting dedicated overloads using pd.Series ?

I think this is pretty complicated, because you need to have overloads for each combination where axis=None and using AnyArrayLike corresponds to Never, but using AnyArrayLike requires axis to be specified and not be None.

I'm going to approve this PR as is - it's an improvement, and if you want to experiment with the overloads, that could be a separate PR. Created #984 to track that.

[Dr-Irv]

thanks @ldouteau