CODEWITHSUNDEEP
pythonlist-comprehension
Nico Schlömer
Nico Schlömer

Reputation: 58821

Function in Python list comprehension, don't eval twice

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?

Upvotes: 2

Views: 73

Answers (2)

buran
buran

Reputation: 14253

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]

Upvotes: 4

Plagon
Plagon

Reputation: 3139

Use the := operator for python >=3.8.

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

Upvotes: 6

Related Questions