From a918530c9adeaac9ffbf34d59c22e76b9101f15f Mon Sep 17 00:00:00 2001 From: "AzureAD\\LouisDouteau" Date: Wed, 21 Aug 2024 09:29:21 +0200 Subject: [PATCH 1/3] Supports arrays in DataFrame.clip (#980) --- pandas-stubs/core/frame.pyi | 4 ++-- tests/test_frame.py | 15 +++++++++++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 8607343c..1ca7c568 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 a8b03e76..b2d417e4 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -527,9 +527,20 @@ 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: + df.clip(lower=lower, upper=upper, axis=axis) def test_types_abs() -> None: From a54526491d0ef8b96ad9bffc58b0dccaa176f0cc Mon Sep 17 00:00:00 2001 From: "AzureAD\\LouisDouteau" Date: Wed, 21 Aug 2024 09:33:08 +0200 Subject: [PATCH 2/3] add missing assert_type (#980) --- tests/test_frame.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index b2d417e4..3f892098 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -540,7 +540,7 @@ def is_none_or_numeric(val: Any) -> bool: with pytest.raises(ValueError): df.clip(lower=lower, upper=upper, axis=axis) else: - df.clip(lower=lower, upper=upper, axis=axis) + assert_type(df.clip(lower=lower, upper=upper, axis=axis), pd.DataFrame) def test_types_abs() -> None: From 6a85b3d4cde316ece27a7a18e4fe4e5c614b698a Mon Sep 17 00:00:00 2001 From: "AzureAD\\LouisDouteau" Date: Wed, 21 Aug 2024 14:13:27 +0200 Subject: [PATCH 3/3] check the content of asserted type --- tests/test_frame.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 3f892098..33aa1616 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -540,7 +540,10 @@ def is_none_or_numeric(val: Any) -> bool: with pytest.raises(ValueError): df.clip(lower=lower, upper=upper, axis=axis) else: - assert_type(df.clip(lower=lower, upper=upper, axis=axis), pd.DataFrame) + check( + assert_type(df.clip(lower=lower, upper=upper, axis=axis), pd.DataFrame), + pd.DataFrame, + ) def test_types_abs() -> None: