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?

Answer 1

Use the := operator for python >=3.8.

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