2

If I do:

def foo():
    yield from range(0,10)
    yield from range(10,20)

for num in foo():
    print(num)

I get a ordered list from 0 to 19. Without change the input of the range functions, is there an easy way to specify that I want a list that goes: 0,10,1,11,2,12...

Basically I first want the first element of every generator. Than I want the second element of every generator and then the third and so on.

Bonus points: Is there a way to change it so that when the generators produce an unequal amount of results, the second generator yields the rest of it's results after the first one is finished?

Martijn Pieters
  • 1,048,767
  • 296
  • 4,058
  • 3,343
Christian
  • 25,249
  • 40
  • 134
  • 225

1 Answers1

8

You are trying to zip() your iterators; do so explicitly:

from itertools import chain

def foo():
    yield from chain.from_iterable(zip(range(10), range(10, 20)))

The use of itertools.chain.from_iterable() lets you continue to use yield from here, flattening out the tuples zip() produces.

Demo:

>>> from itertools import chain
>>> def foo():
...     yield from chain.from_iterable(zip(range(10), range(10, 20)))
... 
>>> list(foo())
[0, 10, 1, 11, 2, 12, 3, 13, 4, 14, 5, 15, 6, 16, 7, 17, 8, 18, 9, 19]

If you have generators of different lengths, you could use itertools.zip_longest():

from itertools import zip_longest

def foo():
    yield from (i for pair in zip_longest(range(10), range(10, 22))
                  for i in pair if i is not None)

I used a different flattening technique here with a double loop in a generator expression.

This all does get tedious, and since you are not using yield from with another generator (so you don't need support for generator.send() and generator.throw() to be propagated), you may as well just make this a proper loop:

def foo():
    for x, y in zip_longest(range(10), range(10, 22)):
        if x is not None:
            yield x
        if y is not None:
            yield y

You can also use the roundrobin() recipe listed in the itertools documentation recipies section:

from itertools import cycle

def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Recipe credited to George Sakkis
    pending = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while pending:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            pending -= 1
            nexts = cycle(islice(nexts, pending))

def foo():
    yield from roundrobin(range(10), range(10, 22))
Martijn Pieters
  • 1,048,767
  • 296
  • 4,058
  • 3,343
  • What if the generators have different length? This solution seems to silently forget the rest of the chain. I would prefer that I still get the rest. – Christian Sep 13 '14 at 14:15
  • 1
    @Christian: then use [`itertools.zip_longest`](https://docs.python.org/3/library/itertools.html#itertools.zip_longest), perhaps filtering out `None` values. But you'll have to do this *explicitly*. – Martijn Pieters Sep 13 '14 at 14:17
  • @MartijnPieters I think I've said this before¹, but... `roundrobin` from the docs, perhaps? ¹http://stackoverflow.com/a/23874310/1763356 – Veedrac Sep 13 '14 at 18:52
  • @Veedrac: I do keep forgetting about roundrobin! – Martijn Pieters Sep 13 '14 at 19:06