MarkS
MarkS

Reputation: 1539

Unpack *args for use in a for loop

Here is my starting point. This works as is.

l1 = [1, 2, 3, 4]
l2 = [5, 6, 7, 8]
l3 = [9, 10, 11, 12]


def interleave(iterable1, iterable2):
    for item1, item2 in zip(iterable1, iterable2):
        yield item1
        yield item2


print(interleave(l1, l2))
print('-' * 30)
print(list(interleave(l1, l2)))

If I want to expand it to use all three lists, I can do this:

l1 = [1, 2, 3, 4]
l2 = [5, 6, 7, 8]
l3 = [9, 10, 11, 12]


def interleave(*args):
    for item1, item2, item3 in zip(*args):
        yield item1
        yield item2
        yield item3


print(interleave(l1, l2, l3))
print('-' * 30)
print(list(interleave(l1, l2, l3)))

However, I 'solved' the issue of receiving any number of input iterables using *args, but my item assignment is still manual.

I want to be able to say:

def interleave(*args):
    for *items in zip(*args):
        yield items

To allow me to unpack any number of input variables into *items, but I get this error: SyntaxError: starred assignment target must be in a list or tuple

I don't want to say *items = <something>. I want to have *items receive any number of input variables.

If I have six lists I don't want to say item1, item2, ..., item6, followed by the same number of yields.

That really isn't very scalable.

Is there a way to do what I am after?

Upvotes: 0

Views: 228

Answers (2)

Olivier Melan&#231;on
Olivier Melan&#231;on

Reputation: 22324

For Python prior to 3.3, use a nested for loop.

def interleave(*args):
    for items in zip(*args):
        for item in items:
            yield item

For Python 3.3+ you can use yield from for generator delegation which is syntactical sugar for the above.

def interleave(*args):
    for items in zip(*args):
        yield from items

Or finally, if you need a more general solution that permits lists of different lengths, use the roundrobin function from the itertools recipes.

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))

Upvotes: 2

Mehdi
Mehdi

Reputation: 4318

A nested loop?

def interleave(*args):
    for items in zip(*args): 
        for item in items:
            yield items

Upvotes: 0

Related Questions