Skip to content

Commit

Permalink
Add support for arrays in lower and upper args of DataFrame.clip (
Browse files Browse the repository at this point in the history
#982)

* Supports arrays in DataFrame.clip (#980)

* add missing assert_type (#980)

* check the content of asserted type
  • Loading branch information
ldouteau authored Aug 21, 2024
1 parent f8b06e1 commit 791777f
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 4 deletions.
4 changes: 2 additions & 2 deletions pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ...,
Expand Down
18 changes: 16 additions & 2 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit 791777f

Please sign in to comment.