diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 8607343c7..1ca7c568e 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1632,8 +1632,8 @@ class DataFrame(NDFrame, OpsMixin): ) -> DataFrame: ... def clip( self, - lower: float | None = ..., - upper: float | None = ..., + lower: float | AnyArrayLike | None = ..., + upper: float | AnyArrayLike | None = ..., *, axis: Axis | None = ..., inplace: _bool = ..., diff --git a/tests/test_frame.py b/tests/test_frame.py index a8b03e76f..33aa1616e 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -527,9 +527,23 @@ def test_types_quantile() -> None: df.quantile(np.array([0.25, 0.75])) -def test_types_clip() -> None: +@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) + df = pd.DataFrame(data={"col1": [20, 12], "col2": [3, 14]}) - df.clip(lower=5, upper=15) + 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, + ) def test_types_abs() -> None: