Rewrite chunk_sequence to work with generators.

This commit is contained in:
voussoir 2020-08-12 11:21:48 -07:00
parent a66aca8a17
commit b15beb73eb

View file

@ -88,26 +88,21 @@ def checkerboard_image(color_1, color_2, image_size, checker_size):
def chunk_sequence(sequence, chunk_length, allow_incomplete=True): def chunk_sequence(sequence, chunk_length, allow_incomplete=True):
''' '''
Given a sequence, divide it into sequences of length `chunk_length`. Given a sequence, yield lists of length `chunk_length`.
allow_incomplete: allow_incomplete:
If True, allow the final chunk to be shorter if the If True, allow the final chunk to be shorter if the
given sequence is not an exact multiple of `chunk_length`. given sequence is not an exact multiple of `chunk_length`.
If False, the incomplete chunk will be discarded. If False, the incomplete chunk will be discarded.
''' '''
(complete, leftover) = divmod(len(sequence), chunk_length) import itertools
if not allow_incomplete: iterator = iter(sequence)
leftover = 0 while True:
chunk = list(itertools.islice(iterator, chunk_length))
chunk_count = complete + min(leftover, 1) if not chunk:
break
chunks = [] if len(chunk) == chunk_length or allow_incomplete:
for x in range(chunk_count): yield chunk
left = chunk_length * x
right = left + chunk_length
chunks.append(sequence[left:right])
return chunks
def comma_space_split(s): def comma_space_split(s):
''' '''