CODEWITHSUNDEEP
pythonpython-3.xlist-comprehension
Steven
Steven

Reputation: 1173

Assignments in python list comprehension

I am looking for a way to do assignments in a list comprehension. I would like to rewrite something like the following piece of code into a list comprehension.

I have this "costly" function:

import time
def f(x):
    time.sleep(1)
    return x + 1

And this loop:

l = []
for value in [1, 2, 3]:
    x = f(value)
    l.append(x + x)

I would like to rewrite this to a list comprehension:

l = [
     f(value) + f(fvalue)
     for value in [1, 2, 3]
]

But since calling f(value) is costly, I would like to minimize the number of calls (this following snippet doesn't run):

l = [
     (x = f(value))
     x + x
     for value in [1, 2, 3]
]

I've read about the assignment expression (:=) (https://www.python.org/dev/peps/pep-0572/#changing-the-scope-rules-for-comprehensions) but I can't seem to figure it out.

Upvotes: 3

Views: 244

Answers (3)

Ferenc Wágner
Ferenc Wágner

Reputation: 266

Here is another way to fake it:

[x+x for value in [1, 2, 3] for x in [f(value)]]

Upvotes: 0

liveware
liveware

Reputation: 72

This can be done with a walrus operator. Note that you need Python 3.8 or later.

l = [x:=10 * value for value in [1, 2, 3]]

With one fast and slow function.

import time

def slow_function(a):
    time.sleep(1)
    return a + 1

def fast_function(b):
    return b + b  


l = [
     fast_function(x := slow_function(value))
     for value in [1, 2, 3]
]

The walrus operator can be skipped in both these examples, but perhaps the real problem would require it.

Upvotes: -1

Lootcifer
Lootcifer

Reputation: 51

My approach would be to nest multiple list comprehension, like

l_new = [x * x * y for x, y in [f(value) for value in [1, 2, 3]]]

So f() should only be called once for each value.

Upvotes: 2

Related Questions