CODEWITHSUNDEEP
pythoniteratorgenerator
Nyps
Nyps

Reputation: 886

Retrieving all iterator values from generator function

Let's say I have a generator function which yields two values:

def gen_func():
    for i in range(5):
        yield i, i**2

I want to retrieve all iterator values of my function. Currently I use this code snippet for that purpose:

x1, x2 = [], []
for a, b in gen_func():
    x1.append(a)
    x2.append(b)

This works for me, but seems a little clunky. Is there a more compact way for coding this? I was thinking something like:

x1, x2 = map(list, zip(*(a, b for a, b in gen_func())))

This, however, just gives me a syntax error.

PS: I am aware that I probably shouldn't use a generator for this purpose, but I need it elsewhere.

Edit: Any type for x1 and x2 would work, however, I prefer list for my case.

Upvotes: 2

Views: 130

Answers (2)

vaultah
vaultah

Reputation: 46603

If x1 and x2 can be tuples, it's sufficient to do

>>> x1, x2 = zip(*gen_func())
>>> x1
(0, 1, 2, 3, 4)
>>> x2
(0, 1, 4, 9, 16)

Otherwise, you could use map to apply list to the iterator:

x1, x2 = map(list, zip(*gen_func()))

Just for fun, the same thing can be done using extended iterable unpacking:

>>> (*x1,), (*x2,) = zip(*gen_func())
>>> x1
[0, 1, 2, 3, 4]
>>> x2
[0, 1, 4, 9, 16]

Upvotes: 6

zwer
zwer

Reputation: 25829

You almost had it:

x1, x2 = map(list, zip(*gen_func()))

Upvotes: 3

Related Questions