CODEWITHSUNDEEP
pythongeneratorenumerate
martineau
martineau

Reputation: 123473

Variable length 'yield'?

I'd like to be able to yield a variable number of items which would allow a generator function something like the following to be written:

x = [1, 2, 3]
y = [4, 5, 6]
z = [7, 8, 9]

def multi_enumerate(*iterables):
    n = 0
    iterators = map(iter, iterables)
    while iterators:
        yield n, *tuple(map(next, iterators))  # invalid syntax
        n += 1

for i,a,b,c in multi_enumerate(x, y, z):
    print i,a,b,c

Anyone know of some way to do this? I know I can yield a tuple, but that would require explicitly unpacking it on the receiving end such as: a,b,c = t[0], t[1], t[2].

Final Solution:

FWIW, here's what I ended up using, based on John Kugelman's excellent answer:

from itertools import izip

def multi_enumerate(*iterables, **kwds):
    start = kwds.get('start')
    if start is None:
        if kwds: raise TypeError(
            "multi_enumerate() accepts only the keyword argument 'start'")
        try:
            iter(iterables[-1])  # last non-keyword arg something iterable?
        except TypeError:        # no, assume it's the start value
            start, iterables = iterables[-1], iterables[:-1]
        else:
            start = 0  # default

    return ((n,)+t for n, t in enumerate(izip(*iterables), start))

The added code is because of my desire to make it also accept an optional non-iterable last argument to specify a starting value other than 0 (or specify it using a start keyword argument), just like the built-in enumerate() function.

Upvotes: 3

Views: 657

Answers (4)

Silas Ray
Silas Ray

Reputation: 26160

This will return an iterator that produces (n, *l[n]) for i n-length ls. It will work directly in place of your multienumerate method.

from itertools import izip
def enumerated_columns(*rows):
    return izip(range(len(rows[0])), *rows)

Upvotes: 0

John Kugelman
John Kugelman

Reputation: 361645

Change the yield statement to this:

yield (n,) + tuple(map(next, iterators))

Or use izip and enumerate to eliminate the whole loop:

from itertools import izip

def multi_enumerate(*iterables):
    return ((n,) + t for n, t in enumerate(izip(*iterables)))

Upvotes: 5

9000
9000

Reputation: 40894

If you want to delegate several items from your generator to some other generator, there's no one-line way to do this — unless you use Python 3.3 that features yield from.

I'd iterate over the nested generators yielding values from them explicitly, something like

yield n
for value in (delegated.next for delegated in iterators):
  yield value

Upvotes: 0

ovgolovin
ovgolovin

Reputation: 13410

I would use chain from itertools:

yield tuple(chain((n,), map(next, iterators)))

Upvotes: 1

Related Questions