Search code examples
pythonpython-3.xpython-itertoolscycle

List comprehension from itertools.cycle generator

My problem was that I needed to deliver batches from an itertools.cycle generator in list form.

A cycle takes an iterable and loops it around itself indefinitely. For example:

>>> my_cycle = itertools.cycle('abc')
>>> next(my_cycle)
'a'
>>> next(my_cycle)
'b'
>>> next(my_cycle)
'c'
>>> next(my_cycle)
'a'

And so on.

The question becomes, how do we deliver a list of batch length n from a cyclic generator, while preserving where we are in the cycle?

Desired output is:

c = itertools.cycle('abc')
batch_size = 2
Out[0]: ['a', 'b']
Out[1]: ['c', 'a']
Out[2]: ['b', 'c']

I am posting my solution in case someone runs into the same problem.

Solution

  • There is an itertools recipe designed for this:

    from itertools import islice, cycle


def take(n, iterable):
    "Return first n items of the iterable as a list"
    return list(islice(iterable, n))


c = cycle("abcdefg")
take(5, c)
# ['a', 'b', 'c', 'd', 'e']