Dargor Dargor - 6 months ago 18
Python Question

Chain repetitions of same generator in python

Consider the following simple generator:

def simple_gen():
for number in range(5):
yield number ** 2


I'm willing to user
itertools.repeat
and
itertools.chain
to chain n times the generator to itself. For clarity consider the following (non-generator) example of the same:

array = [1,2,3,4]
repetitions = itertools.repeat( array ,2)
list(itertools.chain.from_iterable(repetitions)) -> [1,2,3,4,1,2,3,4]


I want the same but using my own generator (simple_gen) instead of array. Of course a simple substitution does not work because
itertools.repeat
repeats the same object and therefore the following repetitions of the generator will be exhausted.

Some ideas of how to achieve this using the itertools module?

I do not want to cast the generator to a list or another container.

Answer

You could convert the generator output to a list first:

repeat(list(simple_gen()), 2)

You can't otherwise repeat a generator output. At most you could recreate the generator count number of times:

from itertools import repeat

def recreate(callable, count=None):
    for c in repeat(callable, count):
        yield from c()

in Python 3 and

from itertools import repeat

def recreate(callable, count=None):
    for c in repeat(callable, count):
        for val in c():
            yield val

for Python 2 and use that instead of chain.from_iterable(repeat(callable(), count)). Note that the generator function is not called, you pass in the function object itself instead.

Demo:

>>> from itertools import repeat
>>> def simple_gen():
...     for number in range(5):
...         yield number ** 2
...
>>> def recreate(callable, count=None):
...     for c in repeat(callable, count):
...         for val in c():
...             yield val
...
>>> list(recreate(simple_gen, 2))
[0, 1, 4, 9, 16, 0, 1, 4, 9, 16]
Comments