Create iterator to return elements from each iterable one by one

Question

I was learning the itertools module and I am trying to make an iterator to return each element from the iterables provided as input.

Agruments   Results
p, q, …     p0, q0, … plast, qlast 

with one more rider that if say the lists are not of the same length then next(it) should return elements from the longer list when the shorter one runs out.

Attempt at solution

import itertools
l1=[1,2,3,4,5,6]
l2=['a','b','c','d']
l=[]
for x,y in itertools.zip_longest(l1,l2):
    l.extend([x,y])
it=iter(x for x in l if x is not None)

Which kind of solves my problem

print(list(it))

Outputs:

[1, 'a', 2, 'b', 3, 'c', 4, 'd', 5, 6]

Is there an easier or better way to do this? I searched for a solution on SO and was not able to get one.

Answer 1

You can use itertools.chain.from_iterable() to flatten the sequence, and use a generator expression to filter out the None values:

from itertools import chain, zip_longest

it = (v for v in chain.from_iterable(zip_longest(l1, l2)) if v is not None)

Rather than use None as the sentinel value, you may want to use a dedicated sentinel so you can use None in the input list:

_sentinel = object()
flattened = chain.from_iterable(zip_longest(l1, l2, fillvalue=_sentinel))
it = (v for v in flattened if v is not _sentinel)

If you want to filter out falsey values, then you can also use filter(None, ...):

it = filter(None, chain.from_iterable(zip_longest(l1, l2)))

Demo:

>>> from itertools import chain, zip_longest
>>> l1 = [1, 2, 3, 4, 5, 6]
>>> l2 = ['a', 'b', 'c', 'd']
>>> it = (v for v in chain.from_iterable(zip_longest(l1, l2)) if v is not None)
>>> list(it)
[1, 'a', 2, 'b', 3, 'c', 4, 'd', 5, 6]

and with a local sentinel:

>>> l1 = [1, None, 2, None, 3, None]
>>> l2 = ['a', 'b', 'c', 'd']
>>> _sentinel = object()
>>> flattened = chain.from_iterable(zip_longest(l1, l2, fillvalue=_sentinel))
>>> it = (v for v in flattened if v is not _sentinel)
>>> list(it)
[1, 'a', None, 'b', 2, 'c', None, 'd', 3, None]

The itertools recipes section also has:

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

Answer 2

If you want a modified version of your code, building a generator from the start (no storing list l):

import itertools
l1=[1,2,3,4,5,6]
l2=['a','b','c','d']

def flat_zip(l1,l2):
    for x,y in itertools.zip_longest(l1,l2):
        if x:
            yield x
        if y:
            yield y
it=flat_zip(l1,l2)

Though I advise for using the builtin solutions above.