diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 259f6be42..cbe882df8 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1738,23 +1738,73 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): @overload def clip( self, - lower: float | AnyArrayLike | None = ..., - upper: float | AnyArrayLike | None = ..., + lower: float | None = ..., + upper: float | None = ..., *, axis: Axis | None = ..., - inplace: Literal[True], + inplace: Literal[False] = ..., **kwargs: Any, - ) -> None: ... + ) -> Self: ... @overload def clip( self, - lower: float | AnyArrayLike | None = ..., - upper: float | AnyArrayLike | None = ..., + lower: AnyArrayLike = ..., + upper: AnyArrayLike | None = ..., *, - axis: Axis | None = ..., + axis: Axis = ..., inplace: Literal[False] = ..., **kwargs: Any, ) -> Self: ... + @overload + def clip( + self, + lower: AnyArrayLike | None = ..., + upper: AnyArrayLike = ..., + *, + axis: Axis = ..., + inplace: Literal[False] = ..., + **kwargs: Any, + ) -> Self: ... + @overload + def clip( # pyright: ignore[reportOverlappingOverload] + self, + lower: None = ..., + upper: None = ..., + *, + axis: Axis | None = ..., + inplace: Literal[True], + **kwargs: Any, + ) -> Self: ... + @overload + def clip( + self, + lower: float | None = ..., + upper: float | None = ..., + *, + axis: Axis | None = ..., + inplace: Literal[True], + **kwargs: Any, + ) -> None: ... + @overload + def clip( + self, + lower: AnyArrayLike = ..., + upper: AnyArrayLike | None = ..., + *, + axis: Axis = ..., + inplace: Literal[True], + **kwargs: Any, + ) -> None: ... + @overload + def clip( + self, + lower: AnyArrayLike | None = ..., + upper: AnyArrayLike = ..., + *, + axis: Axis = ..., + inplace: Literal[True], + **kwargs: Any, + ) -> None: ... def copy(self, deep: _bool = ...) -> Self: ... def cummax( self, axis: Axis | None = ..., skipna: _bool = ..., *args: Any, **kwargs: Any diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 3d9a73301..d7e77afee 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -1394,6 +1394,16 @@ class Series(IndexOpsMixin[S1], NDFrame): subset: _str | Sequence[_str] | None = ..., ) -> Scalar | Series[S1]: ... @overload + def clip( # pyright: ignore[reportOverlappingOverload] + self, + lower: None = ..., + upper: None = ..., + *, + axis: AxisIndex | None = ..., + inplace: Literal[True], + **kwargs: Any, + ) -> Self: ... + @overload def clip( self, lower: AnyArrayLike | float | None = ..., diff --git a/tests/test_frame.py b/tests/test_frame.py index 47dbc6460..f7aa071e3 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -659,23 +659,218 @@ def test_types_quantile() -> None: df.quantile(np.array([0.25, 0.75])) -@pytest.mark.parametrize("lower", [None, 5, pd.Series([3, 4])]) -@pytest.mark.parametrize("upper", [None, 15, pd.Series([12, 13])]) -@pytest.mark.parametrize("axis", [None, 0, "index"]) -def test_types_clip(lower, upper, axis) -> None: - def is_none_or_numeric(val: Any) -> bool: - return val is None or isinstance(val, int | float) - +def test_dataframe_clip() -> None: + """Test different clipping combinations for dataframe.""" df = pd.DataFrame(data={"col1": [20, 12], "col2": [3, 14]}) - uses_array = not (is_none_or_numeric(lower) and is_none_or_numeric(upper)) - if uses_array and axis is None: - with pytest.raises(ValueError): - df.clip(lower=lower, upper=upper, axis=axis) - else: - check( - assert_type(df.clip(lower=lower, upper=upper, axis=axis), pd.DataFrame), - pd.DataFrame, - ) + if TYPE_CHECKING_INVALID_USAGE: + df.clip(lower=pd.Series([4, 5]), upper=None, axis=None) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType] + df.clip(lower=None, upper=pd.Series([4, 5]), axis=None) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType] + df.clip(lower=pd.Series([1, 2]), upper=pd.Series([4, 5]), axis=None) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType] + df.copy().clip(lower=pd.Series([1, 2]), upper=None, axis=None, inplace=True) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType] + df.copy().clip(lower=None, upper=pd.Series([1, 2]), axis=None, inplace=True) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType] + df.copy().clip(lower=pd.Series([4, 5]), upper=pd.Series([1, 2]), axis=None, inplace=True) # type: ignore[call-overload] # pyright: ignore[reportCallIssue, reportArgumentType] + + check( + assert_type(df.clip(lower=None, upper=None, axis=None), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=5, upper=None, axis=None), pd.DataFrame), pd.DataFrame + ) + check( + assert_type(df.clip(lower=None, upper=15, axis=None), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type( + df.clip(lower=None, upper=None, axis=None, inplace=True), pd.DataFrame + ), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=5, upper=None, axis=None, inplace=True), None), + type(None), + ) + check( + assert_type(df.clip(lower=None, upper=15, axis=None, inplace=True), None), + type(None), + ) + + check( + assert_type(df.clip(lower=None, upper=None, axis=0), pd.DataFrame), pd.DataFrame + ) + check(assert_type(df.clip(lower=5, upper=None, axis=0), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.clip(lower=None, upper=15, axis=0), pd.DataFrame), pd.DataFrame + ) + check( + assert_type(df.clip(lower=pd.Series([1, 2]), upper=None, axis=0), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=None, upper=pd.Series([1, 2]), axis=0), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type( + df.clip(lower=None, upper=None, axis="index", inplace=True), pd.DataFrame + ), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=5, upper=None, axis="index", inplace=True), None), + type(None), + ) + check( + assert_type(df.clip(lower=None, upper=15, axis="index", inplace=True), None), + type(None), + ) + check( + assert_type( + df.clip(lower=pd.Series([1, 2]), upper=None, axis="index", inplace=True), + None, + ), + type(None), + ) + check( + assert_type( + df.clip(lower=None, upper=pd.Series([1, 2]), axis="index", inplace=True), + None, + ), + type(None), + ) + check( + assert_type(df.clip(lower=None, upper=None, axis="index"), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=5, upper=None, axis="index"), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=None, upper=15, axis="index"), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type( + df.clip(lower=pd.Series([1, 2]), upper=None, axis="index"), pd.DataFrame + ), + pd.DataFrame, + ) + check( + assert_type( + df.clip(lower=None, upper=pd.Series([1, 2]), axis="index"), pd.DataFrame + ), + pd.DataFrame, + ) + check( + assert_type( + df.clip(lower=None, upper=None, axis=0, inplace=True), pd.DataFrame + ), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=5, upper=None, axis=0, inplace=True), None), + type(None), + ) + check( + assert_type(df.clip(lower=None, upper=15, axis=0, inplace=True), None), + type(None), + ) + + # without lower + check(assert_type(df.clip(upper=None, axis=None), pd.DataFrame), pd.DataFrame) + check(assert_type(df.clip(upper=15, axis=None), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.clip(upper=None, axis=None, inplace=True), pd.DataFrame), + pd.DataFrame, + ) + check(assert_type(df.clip(upper=15, axis=None, inplace=True), None), type(None)) + + check(assert_type(df.clip(upper=None, axis=0), pd.DataFrame), pd.DataFrame) + check(assert_type(df.clip(upper=15, axis=0), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.clip(upper=pd.Series([1, 2]), axis=0), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(upper=None, axis="index", inplace=True), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(upper=15, axis="index", inplace=True), None), + type(None), + ) + check( + assert_type(df.clip(upper=pd.Series([1, 2]), axis="index", inplace=True), None), + type(None), + ) + check(assert_type(df.clip(upper=None, axis="index"), pd.DataFrame), pd.DataFrame) + check(assert_type(df.clip(upper=15, axis="index"), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.clip(upper=pd.Series([1, 2]), axis="index"), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(upper=None, axis=0, inplace=True), pd.DataFrame), + pd.DataFrame, + ) + check(assert_type(df.clip(upper=15, axis=0, inplace=True), None), type(None)) + + # without upper + check( + assert_type(df.clip(lower=None, axis=None), pd.DataFrame), + pd.DataFrame, + ) + check(assert_type(df.clip(lower=5, axis=None), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.clip(lower=None, axis=None, inplace=True), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=5, axis=None, inplace=True), None), + type(None), + ) + check( + assert_type(df.clip(lower=None, axis=None, inplace=True), pd.DataFrame), + pd.DataFrame, + ) + + check(assert_type(df.clip(lower=None, axis=0), pd.DataFrame), pd.DataFrame) + check(assert_type(df.clip(lower=5, axis=0), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.clip(lower=pd.Series([1, 2]), axis=0), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=None, axis="index", inplace=True), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=5, axis="index", inplace=True), None), + type(None), + ) + check( + assert_type(df.clip(lower=pd.Series([1, 2]), axis="index", inplace=True), None), + type(None), + ) + check( + assert_type(df.clip(lower=None, axis="index"), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=5, axis="index"), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=pd.Series([1, 2]), axis="index"), pd.DataFrame), + pd.DataFrame, + ) + check( + assert_type(df.clip(lower=None, axis=0, inplace=True), pd.DataFrame), + pd.DataFrame, + ) + check(assert_type(df.clip(lower=5, axis=0, inplace=True), None), type(None)) def test_types_abs() -> None: diff --git a/tests/test_series.py b/tests/test_series.py index f84648df1..265ac7711 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -582,9 +582,130 @@ def test_types_quantile() -> None: def test_types_clip() -> None: + """Test different overloads of Series.clip GH984.""" s = pd.Series([-10, 2, 3, 10]) - s.clip(lower=0, upper=5) - s.clip(lower=0, upper=5, inplace=True) + lower = pd.Series([-10, -10, -3, -10]) + upper = pd.Series([50, 52, 53, 51]) + check( + assert_type(s.clip(lower=None, upper=None), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check( + assert_type(s.clip(lower=0, upper=5), "pd.Series[int]"), pd.Series, np.integer + ) + check( + assert_type(s.clip(lower=0, upper=None), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check( + assert_type(s.clip(lower=None, upper=5), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check( + assert_type(s.clip(lower=None, upper=None, inplace=True), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check(assert_type(s.clip(lower=0, upper=5, inplace=True), None), type(None)) + check(assert_type(s.clip(lower=0, upper=None, inplace=True), None), type(None)) + check( + assert_type(s.clip(lower=None, upper=None, inplace=True), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check( + assert_type(s.clip(lower=None, upper=5, inplace=True), None), + type(None), + ) + check( + assert_type(s.clip(lower=lower, upper=upper), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check( + assert_type(s.clip(lower=lower, upper=upper, axis=0), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check( + assert_type(s.clip(lower=lower, upper=upper, axis="index"), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check(assert_type(s.clip(lower=lower, upper=upper, inplace=True), None), type(None)) + check( + assert_type(s.clip(lower=lower, upper=upper, axis=0, inplace=True), None), + type(None), + ) + check( + assert_type(s.clip(lower=lower, upper=upper, axis="index", inplace=True), None), + type(None), + ) + + # without lower + check(assert_type(s.clip(upper=None), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.clip(upper=5), "pd.Series[int]"), pd.Series, np.integer) + check( + assert_type(s.clip(upper=None, inplace=True), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check(assert_type(s.clip(upper=5, inplace=True), None), type(None)) + check(assert_type(s.clip(upper=upper), "pd.Series[int]"), pd.Series, np.integer) + check( + assert_type(s.clip(upper=upper, axis=0), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check( + assert_type(s.clip(upper=upper, axis="index"), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check(assert_type(s.clip(upper=upper, inplace=True), None), type(None)) + check(assert_type(s.clip(upper=upper, axis=0, inplace=True), None), type(None)) + check( + assert_type(s.clip(upper=upper, axis="index", inplace=True), None), type(None) + ) + + # without upper + check(assert_type(s.clip(lower=None), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.clip(lower=0), "pd.Series[int]"), pd.Series, np.integer) + check(assert_type(s.clip(lower=None), "pd.Series[int]"), pd.Series, np.integer) + check( + assert_type(s.clip(lower=None, inplace=True), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check(assert_type(s.clip(lower=0, inplace=True), None), type(None)) + check( + assert_type(s.clip(lower=None, inplace=True), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check(assert_type(s.clip(lower=lower), "pd.Series[int]"), pd.Series, np.integer) + check( + assert_type(s.clip(lower=lower, axis=0), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check( + assert_type(s.clip(lower=lower, axis="index"), "pd.Series[int]"), + pd.Series, + np.integer, + ) + check(assert_type(s.clip(lower=lower, inplace=True), None), type(None)) + check(assert_type(s.clip(lower=lower, axis=0, inplace=True), None), type(None)) + check( + assert_type(s.clip(lower=lower, axis="index", inplace=True), None), type(None) + ) + + if TYPE_CHECKING_INVALID_USAGE: + s.clip(lower=lower, axis=1) # type: ignore[call-overload] # pyright: ignore[reportArgumentType] + s.clip(lower=lower, axis="column") # type: ignore[call-overload] # pyright: ignore[reportArgumentType] def test_types_abs() -> None: