Search code examples
pythongeneratorpython-itertools

Cycle through nested generators once and repeat


I want to yield through 2 different itertools.count. I have combined the two generators using itertools.chain.from_iterable

This is the code I have written for it.

return itertools.chain.from_iterable([itertools.count(start=2, step=2), itertools.count(start=7, step=7)])

The problem is that it is trying to finish the first counter (step 2) before proccding to yield over next counter (step 7)

Output from the above sample code:

2
4
6
8
10
...

But I want to cycle over alternatively.

Expected Output:

2  # 2*1
7  # 7*1
4  # 2*2
14 # 7*2
6  # 2*3
21 # 7*3
8  # 2*4
28 # 7*4

Here are the other ways I have tried so far:

yield from [elem for elem in [next(count(start=2, step=2)), next(count(start=7, step=7))]]

The above cycles alternatively but the counter resets after each yield.

Output from the above code sample:

2
7
2
7
2
7

I want this to be implemented entirely on itertools or list comprehension since they are memory optimized, hence I expect the function to return a generator object. Also, it would be better if the solution is on a single line.

EDIT:

As suggested by jonrsharpe in the comment, I have implemented roundrobin iter technique and I am able to fetch the desired output.

from itertools import count, cycle

def pattern_generator():
    return cycle(iter(it).__next__ for it in [
        count(start=2, step=2),
        count(start=7, step=7),
    ])

gen = pattern_generator()

print(next(gen)())
print(next(gen)())
print(next(gen)())
print(next(gen)())

I am satisfied with this output. But, is it possible to call next without calling the iter's next method? i.e., without using () in next(gen)()?

Thanks in advance.


Solution

  • You can make generator in various ways

    inline

    #for i in (i for t in zip('abc',range(3)) for i in t):
    #EDIT: more readable solution
    for i in itertools.chain.from_iterable(zip('abc',range(3))):
        print(i)
    

    EDIT 2: explanation
    zip connects n-th of each iterable (returns sequence of tuples ('a', 0), ('b', 1) ...)
    so this roughly translates to itertools.chain.from_iterable([('a', 0), ('b', 1), ...])
    calling chain.from_iterable is similar to calling chain
    so now we have chain(('a', 0), ('b', 1), ('c', 2))
    since tuples are iterables, chain iterates through ('a', 0), and then ('b', 1) and so on
    from_iterable and zip are both needed because neither actually creates list [('a', 0), ('b', 1), ...] (which in your case would be infinite)

    function 1

    def alternate(*iterables):
        for t in zip(*iterables):
            yield from t # or for i in t: yield i
    

    function 2

    def alternate(*iterables):
        iterables = [iter(it) for it in iterables]
        while True:
            try:
                for it in iterables:
                    yield next(it)
            except StopIteration:
                break
    

    result

    for i in alternate('abc', range(3)):
        print(i)
    a
    0
    b
    1
    c
    2
    

    function approaches also give larger flexibility