From b1e35f660473a2e2c80817687da0359bfa237054 Mon Sep 17 00:00:00 2001 From: Ethan Dalool Date: Tue, 11 Aug 2020 23:12:22 -0700 Subject: [PATCH] Rewrite chunk_sequence to work with generators. At the cost of always returning lists instead of the input's type, which was never relevant anyway. --- ycdl/helpers.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/ycdl/helpers.py b/ycdl/helpers.py index 19d7ab2..925d342 100644 --- a/ycdl/helpers.py +++ b/ycdl/helpers.py @@ -1,24 +1,20 @@ 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`. - :param allow_incomplete: If True, allow the final chunk to be shorter if the + allow_incomplete: + If True, allow the final chunk to be shorter if the given sequence is not an exact multiple of `chunk_length`. If False, the incomplete chunk will be discarded. ''' - (complete, leftover) = divmod(len(sequence), chunk_length) - if not allow_incomplete: - leftover = 0 - - chunk_count = complete + min(leftover, 1) - - chunks = [] - for x in range(chunk_count): - left = chunk_length * x - right = left + chunk_length - chunks.append(sequence[left:right]) - - return chunks + import itertools + iterator = iter(sequence) + while True: + chunk = list(itertools.islice(iterator, chunk_length)) + if not chunk: + break + if len(chunk) == chunk_length or allow_incomplete: + yield chunk def truthystring(s): if isinstance(s, (bool, int)) or s is None: