Skip to content

Commit

Permalink
fix(python): Raise if pass a negative n into clear
Browse files Browse the repository at this point in the history
  • Loading branch information
reswqa committed Apr 2, 2024
1 parent 758b55a commit bab7462
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 8 deletions.
17 changes: 9 additions & 8 deletions py-polars/polars/dataframe/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -6912,17 +6912,18 @@ def clear(self, n: int = 0) -> Self:
│ null ┆ null ┆ null │
└──────┴──────┴──────┘
"""
if n < 0:
msg = "n should be greater than or equal to 0."
raise ValueError(msg)
# faster path
if n == 0:
return self._from_pydf(self._df.clear())
if n > 0 or len(self) > 0:
return self.__class__(
{
nm: pl.Series(name=nm, dtype=tp).extend_constant(None, n)
for nm, tp in self.schema.items()
}
)
return self.clone()
return self.__class__(
{
nm: pl.Series(name=nm, dtype=tp).extend_constant(None, n)
for nm, tp in self.schema.items()
}
)

def clone(self) -> Self:
"""
Expand Down
4 changes: 4 additions & 0 deletions py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -4759,6 +4759,10 @@ def clear(self, n: int = 0) -> Series:
null
]
"""
if n < 0:
msg = "n should be greater than or equal to 0."
raise ValueError(msg)
# faster path
if n == 0:
return self._from_pyseries(self._s.clear())
s = (
Expand Down
10 changes: 10 additions & 0 deletions py-polars/tests/unit/operations/test_clear.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,3 +73,13 @@ def test_clear_series_object_starting_with_null() -> None:
assert result.dtype == s.dtype
assert result.name == s.name
assert result.is_empty()


def test_clear_raise_negative_n() -> None:
s = pl.Series([1, 2, 3])
with pytest.raises(ValueError, match="n should be greater than or equal to 0"):
s.clear(-1)

df = pl.DataFrame({"a": [1, 2, 3]})
with pytest.raises(ValueError, match="n should be greater than or equal to 0"):
df.clear(-1)

0 comments on commit bab7462

Please sign in to comment.