Function in Python list comprehension, don't eval twice

Question

I’m composing a Python list from an input list run through a transforming function. I would like to include only those items in the output list for which the result isn’t None. This works:

def transform(n):
    # expensive irl, so don't execute twice
    return None if n == 2 else n**2


a = [1, 2, 3]

lst = []
for n in a:
    t = transform(n)
    if t is not None:
        lst.append(t)

print(lst)

[1, 9]

I have a hunch that this can be simplified with a comprehension. However, the straighforward solution

def transform(n):
    return None if n == 2 else n**2


a = [1, 2, 3]
lst = [transform(n) for n in a if transform(n) is not None]

print(lst)

is no good since transform() is applied twice to each entry. Any way around this?

Asked By: Nico Schlömer

||

Source

Answer 1

Use the := operator for python >=3.8.

lst = [t for n in a if (t:= transform(n)) is not None]

Answered By: Plagon

Answer 2

If not able/don’t want to use walrus operator, one can use @functools.lru_cache to cache the result from calling the function and avoid calling it twice

import functools

eggs = [2, 4, 5, 3, 2]

@functools.lru_cache
def spam(foo):
    print(foo) # to demonstrate each call
    return None if foo % 2 else foo

print([spam(n) for n in eggs if spam(n) is not None])

output

2
4
5
3
[2, 4, 2]

Compared with walrus operator (currently accepted answer) this will be the better option if there are duplicate values in the input list, i.e. walrus operator will always run the function once per element in the input list. Note, you may combine finctools.lru_cache with walrus operator, e.g. for readability.

eggs = [2, 4, 5, 3, 2]

def spam(foo):
    print(foo) # to demonstrate each call
    return None if foo % 2 else foo

print([bar for n in eggs if (bar:=spam(n)) is not None])

output

2
4
5
3
2
[2, 4, 2]

Answered By: buran

Function in Python list comprehension, don't eval twice

Question:

Answers: