Skip to content

Commit defe68c

Browse files
committed
Raise in partition_all when length is invalid
Sometimes an iterable's `__len__` is incorrect, which can lead to bad output from `partition_all`. We now raise an exception (when we previously output bad data) in these cases so the user can more easily find the invalid iterable and fix it. Fixes #602.
1 parent 08f2604 commit defe68c

File tree

2 files changed

+20
-1
lines changed

2 files changed

+20
-1
lines changed

toolz/itertoolz.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -732,7 +732,12 @@ def partition_all(n, seq):
732732
try:
733733
# If seq defines __len__, then
734734
# we can quickly calculate where no_pad starts
735-
yield prev[:len(seq) % n]
735+
end = len(seq) % n
736+
if prev[end - 1] is no_pad or prev[end] is not no_pad:
737+
raise LookupError(
738+
'The sequence passed to `parition_all` has invalid length'
739+
)
740+
yield prev[:end]
736741
except TypeError:
737742
# Get first index of no_pad without using .index()
738743
# https://github.com/pytoolz/toolz/issues/387

toolz/tests/test_itertoolz.py

+14
Original file line numberDiff line numberDiff line change
@@ -355,6 +355,20 @@ def __eq__(self, other):
355355
assert list(partition_all(4, [obj]*7)) == result
356356
assert list(partition_all(4, iter([obj]*7))) == result
357357

358+
# Test invalid __len__: https://github.com/pytoolz/toolz/issues/602
359+
class ListWithBadLength(list):
360+
def __init__(self, contents, off_by=1):
361+
self.off_by = off_by
362+
super().__init__(contents)
363+
364+
def __len__(self):
365+
return super().__len__() + self.off_by
366+
367+
too_long_list = ListWithBadLength([1, 2], off_by=+1)
368+
assert raises(LookupError, lambda: list(partition_all(5, too_long_list)))
369+
too_short_list = ListWithBadLength([1, 2], off_by=-1)
370+
assert raises(LookupError, lambda: list(partition_all(5, too_short_list)))
371+
358372

359373
def test_count():
360374
assert count((1, 2, 3)) == 3

0 commit comments

Comments
 (0)