3

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.

Steven
  • 1,123
  • 5
  • 14
  • 31

3 Answers3

2

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.

Lootcifer
  • 51
  • 2
  • Thank you for your answer. The problem for me would be that we introduce an extra loop which makes it a bit unreadable to me – Steven Jan 16 '22 at 17:52
0

Here is another way to fake it:

[x+x for value in [1, 2, 3] for x in [f(value)]]
Ferenc Wágner
  • 246
  • 2
  • 9
-1

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.

liveware
  • 72
  • 5
  • I think you don’t even need walrus assignment here, right? – rv.kvetch Jan 16 '22 at 13:15
  • Yeah for the original example I made this is true. I updated the post to demonstrate my problem more clearly. – Steven Jan 16 '22 at 13:22
  • In that case you could write one fast and one slow function, like this. – liveware Jan 16 '22 at 16:14
  • What would be the use of the walrus operator? Since `x` is not (re-)used – Steven Jan 16 '22 at 16:23
  • The original question looks like a simplification of a more complex real problem. Showing that a variable can be assigned inside a list comprehension might help someone solving an issue. I explain at the bottom that the walrus operator is not needed in this specific example. – liveware Jan 16 '22 at 16:28
  • I see. Is there a way to do it without creating extra functions? – Steven Jan 16 '22 at 17:49
  • 1
    You could write it as e.g. `l = [(x := f(value), x + x)[1] for value in [1, 2, 3]]`, if you're willing to create a tuple and take the second element of it inside the list comprehension. – liveware Jan 16 '22 at 19:06
  • Thank you, that's seems the best option. I think Python should add support for (multiple) assignments within list comprehensions, since using list comprehensions over for-loops is considered to be pythonic. – Steven Jan 26 '22 at 08:59