Function in Python list comprehension, don't eval twice-CodePudding

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?

CodePudding user response：

Use the := operator for python >=3.8.

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

CodePudding user response：

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 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.

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]