Skip to content

Instantly share code, notes, and snippets.

@cypreess
Last active December 18, 2015 21:48
Show Gist options
  • Save cypreess/5849702 to your computer and use it in GitHub Desktop.
Save cypreess/5849702 to your computer and use it in GitHub Desktop.
Reading from iterator N elements by a time - the python way.
>>> from itertools import *
# First try is like that - OK it works
>>> count = 5
>>> iterator = xrange(100)
>>> list(imap(None, *([iter(iterator)] * count)))
[(0, 1, 2, 3, 4), (5, 6, 7, 8, 9), (10, 11, 12, 13, 14), (15, 16, 17, 18, 19), (20, 21, 22, 23, 24), (25, 26, 27, 28, 29), (30, 31, 32, 33, 34), (35, 36, 37, 38, 39), (40, 41, 42, 43, 44), (45, 46, 47, 48, 49), (50, 51, 52, 53, 54), (55, 56, 57, 58, 59), (60, 61, 62, 63, 64), (65, 66, 67, 68, 69), (70, 71, 72, 73, 74), (75, 76, 77, 78, 79), (80, 81, 82, 83, 84), (85, 86, 87, 88, 89), (90, 91, 92, 93, 94), (95, 96, 97, 98, 99)]
# Apparently imap(None,...) is equivalent to izip() - so we can simplify it
>>> count = 5
>>> iterator = xrange(100)
>>> list(izip(*([iter(iterator)] * count)))
[(0, 1, 2, 3, 4), (5, 6, 7, 8, 9), (10, 11, 12, 13, 14), (15, 16, 17, 18, 19), (20, 21, 22, 23, 24), (25, 26, 27, 28, 29), (30, 31, 32, 33, 34), (35, 36, 37, 38, 39), (40, 41, 42, 43, 44), (45, 46, 47, 48, 49), (50, 51, 52, 53, 54), (55, 56, 57, 58, 59), (60, 61, 62, 63, 64), (65, 66, 67, 68, 69), (70, 71, 72, 73, 74), (75, 76, 77, 78, 79), (80, 81, 82, 83, 84), (85, 86, 87, 88, 89), (90, 91, 92, 93, 94), (95, 96, 97, 98, 99)]
# But it does not works for every length, you see here not all elements are returned
>>> count = 6
>>> iterator = xrange(100)
>>> list(izip(*([iter(iterator)] * count)))
[(0, 1, 2, 3, 4, 5), (6, 7, 8, 9, 10, 11), (12, 13, 14, 15, 16, 17), (18, 19, 20, 21, 22, 23), (24, 25, 26, 27, 28, 29), (30, 31, 32, 33, 34, 35), (36, 37, 38, 39, 40, 41), (42, 43, 44, 45, 46, 47), (48, 49, 50, 51, 52, 53), (54, 55, 56, 57, 58, 59), (60, 61, 62, 63, 64, 65), (66, 67, 68, 69, 70, 71), (72, 73, 74, 75, 76, 77), (78, 79, 80, 81, 82, 83), (84, 85, 86, 87, 88, 89), (90, 91, 92, 93, 94, 95)]
# But you can use izip_longest
>>> count = 6
>>> iterator = xrange(100)
>>> list(izip_longest( *([iter(iterator)] * count)))
[(0, 1, 2, 3, 4, 5), (6, 7, 8, 9, 10, 11), (12, 13, 14, 15, 16, 17), (18, 19, 20, 21, 22, 23), (24, 25, 26, 27, 28, 29), (30, 31, 32, 33, 34, 35), (36, 37, 38, 39, 40, 41), (42, 43, 44, 45, 46, 47), (48, 49, 50, 51, 52, 53), (54, 55, 56, 57, 58, 59), (60, 61, 62, 63, 64, 65), (66, 67, 68, 69, 70, 71), (72, 73, 74, 75, 76, 77), (78, 79, 80, 81, 82, 83), (84, 85, 86, 87, 88, 89), (90, 91, 92, 93, 94, 95), (96, 97, 98, 99, None, None)]
# Third "but", the last element has None, we would like to get rid of them
>>> count = 6
>>> iterator = xrange(100)
>>> list(imap(lambda t: filter(lambda v: v is not None, t), itertools.izip_longest(*([iter(iterator)] * count))))
[(0, 1, 2, 3, 4, 5), (6, 7, 8, 9, 10, 11), (12, 13, 14, 15, 16, 17), (18, 19, 20, 21, 22, 23), (24, 25, 26, 27, 28, 29), (30, 31, 32, 33, 34, 35), (36, 37, 38, 39, 40, 41), (42, 43, 44, 45, 46, 47), (48, 49, 50, 51, 52, 53), (54, 55, 56, 57, 58, 59), (60, 61, 62, 63, 64, 65), (66, 67, 68, 69, 70, 71), (72, 73, 74, 75, 76, 77), (78, 79, 80, 81, 82, 83), (84, 85, 86, 87, 88, 89), (90, 91, 92, 93, 94, 95), (96, 97, 98, 99)]
@lukaszb
Copy link

lukaszb commented Jun 26, 2013

With some tests:

import unittest
import itertools


def ichunks(seq, n):
    """
    Returns generator of lists being n-element chunks of given sequence.
    """
    seq = iter(seq) # make sure we have generator like sequence
    get_next = lambda: list(itertools.islice(seq, n))
    for items in iter(get_next, []):
        yield items


def gen_range(start, stop):
    while start <= stop:
        yield start
        start += 1


class TestGenRange(unittest.TestCase):

    def test_simple(self):
        self.assertEqual(list(gen_range(1, 3)), [1, 2, 3])
        self.assertEqual(list(gen_range(0, 2)), [0, 1, 2])
        self.assertEqual(list(gen_range(3, 3)), [3])
        self.assertEqual(list(gen_range(3, 4)), [3, 4])


class TestChunks(unittest.TestCase):

    def assertChunksAreFine(self, sequence, n, expected):
        self.assertEqual(list(ichunks(sequence, n)), list(expected))

    def test_simple(self):
        self.assertChunksAreFine(gen_range(1, 6), 3, [[1, 2, 3], [4, 5, 6]])

    def test_with_orphans(self):
        self.assertChunksAreFine(gen_range(1, 5), 2, [[1, 2], [3, 4], [5]])

    def test_list_with_orphans(self):
        self.assertChunksAreFine([1, 2, 3, 4, 5], 3, [
            [1, 2, 3], [4, 5]])

    def test_string_with_orphans(self):
        self.assertChunksAreFine('abcdefghij', 3, [
            ['a', 'b', 'c'], ['d', 'e', 'f'], ['g', 'h', 'i'], ['j']])


if __name__ == '__main__':

    unittest.main()

gen_range is defined for Python 3 (which doesn't have xrange). Besides, xrange is special generator, in example it knows it's length.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment