Skip to content

Commit

Permalink
.skip/.truncate: make no-op if arguments are None
Browse files Browse the repository at this point in the history
  • Loading branch information
ebonnal committed Jan 21, 2025
1 parent f473018 commit 101427f
Show file tree
Hide file tree
Showing 4 changed files with 20 additions and 31 deletions.
4 changes: 2 additions & 2 deletions streamable/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,10 +41,10 @@
validate_group_interval,
validate_group_size,
validate_iterator,
validate_optional_count,
validate_skip_args,
validate_throttle_interval,
validate_throttle_per_period,
validate_truncate_args,
)

with suppress(ImportError):
Expand Down Expand Up @@ -210,7 +210,7 @@ def truncate(
when: Optional[Callable[[T], Any]] = None,
) -> Iterator[T]:
validate_iterator(iterator)
validate_truncate_args(count, when)
validate_optional_count(count)
if count is not None:
iterator = CountTruncateIterator(iterator, count)
if when is not None:
Expand Down
4 changes: 2 additions & 2 deletions streamable/stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,10 @@
validate_concurrency,
validate_group_interval,
validate_group_size,
validate_optional_count,
validate_skip_args,
validate_throttle_interval,
validate_throttle_per_period,
validate_truncate_args,
validate_via,
)

Expand Down Expand Up @@ -489,7 +489,7 @@ def truncate(
Returns:
Stream[T]: A stream of at most `count` upstream elements not satisfying the `when` predicate.
"""
validate_truncate_args(count, when)
validate_optional_count(count)
return TruncateStream(self, count, when)


Expand Down
19 changes: 4 additions & 15 deletions streamable/util/validationtools.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,9 @@ def validate_count(count: int):
if count >= sys.maxsize:
raise ValueError(f"`count` must be < sys.maxsize but got {count}")

def validate_optional_count(count: Optional[int]):
if count is not None:
validate_count(count)

def validate_throttle_per_period(per_period_arg_name: str, value: int) -> None:
if value < 1:
Expand All @@ -56,24 +59,10 @@ def validate_throttle_interval(interval: datetime.timedelta) -> None:
if interval < datetime.timedelta(0):
raise ValueError(f"`interval` must be >= 0 but got {repr(interval)}")


def validate_truncate_args(
count: Optional[int] = None, when: Optional[Callable[[T], Any]] = None
) -> None:
if count is None:
if when is None:
raise ValueError("`count` and `when` cannot both be None")
else:
validate_count(count)


def validate_skip_args(
count: Optional[int] = None, until: Optional[Callable[[T], Any]] = None
) -> None:
if count is None:
if until is None:
raise ValueError("`count` and `until` cannot both be None")
else:
if count is not None:
if until is not None:
raise ValueError("`count` and `until` cannot both be set")
validate_count(count)
24 changes: 12 additions & 12 deletions tests/test_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -820,12 +820,11 @@ def test_skip(self) -> None:
):
Stream(src).skip(0, until=bool)

with self.assertRaisesRegex(
ValueError,
"`count` and `until` cannot both be None",
msg="`skip` must raise ValueError if both `count` and `until` are None",
):
Stream(src).skip()
self.assertListEqual(
list(Stream(src).skip()),
list(src),
msg="`skip` must be no-op if both `count` and `until` are None",
)

for count in [0, 1, 3]:
self.assertListEqual(
Expand Down Expand Up @@ -857,17 +856,18 @@ def test_skip(self) -> None:
)

def test_truncate(self) -> None:
with self.assertRaisesRegex(
ValueError,
"`count` and `when` cannot both be None",
):
Stream(src).truncate()

self.assertListEqual(
list(Stream(src).truncate(N * 2)),
list(src),
msg="`truncate` must be ok with count >= stream length",
)

self.assertListEqual(
list(Stream(src).truncate()),
list(src),
msg="`truncate must be no-op if both `count` and `when` are None",
)

self.assertListEqual(
list(Stream(src).truncate(2)),
[0, 1],
Expand Down

0 comments on commit 101427f

Please sign in to comment.