Unpacking return value from function in List Comprehension

Question

This code:

def f(v):
    return (v,v+1)

a = [f(i) for i in range(3)]

gives:

[(0, 1), (1, 2), (2, 3)]

I wish to change the comprehension so that it gives:

[0, 1, 1, 2, 2, 3]

How do I do this?

Answer 1

In a list comprehension you can use multiple for ... in ...:

>>> def f(v):
...     return (v,v+1)
...
>>> [x for i in range(3) for x in f(i)]
[0, 1, 1, 2, 2, 3]

Or, you can use itertools.chain.from_iterable:

>>> import itertools
>>> list(itertools.chain.from_iterable(f(i) for i in range(3)))
[0, 1, 1, 2, 2, 3]