pythonpython-assignment-expression

Use Output from expression part of list comprehension in conditional


I want to call a method and add the returned value to a list only if the value is not None, in place.

Approach 1

list = []
for item in items:
    temp = func(item)
    if temp is not None:
        list.append(temp)

Approach 2

 [func(item) for item in items if func(item) is not None]

The first approach is more verbose than I want but accomplishes my goal. The second approach is in place, but requires two calls to func(item). If possible I would like to reuse the first call in a var that I then use in the comparison.

Not proper python but in the vein of what I am looking for

 [(a = func(item)) for item in items if a is not None]

Solution

  • This is what the assignment expression, aka the walrus operator, :=, introduced in Python 3.8, supports:

    [a for item in items if (a := func(item)) is not None]
    

    The test precedes the production of a value, so you use the walrus in the condition, not the production expression.

    Pre-3.8, you're stuck wrapping a listcomp around a map or genexpr:

    [a for a in map(func, items) if a is not None]  # wrapping map
    [a for a in (func(item) for item in items) if a is not None]  # wrapping genexpr
    

    Both of the above are equivalent; the former might be slightly more performant (especially if func is a built-in implemented in C), but it's mostly personal preference (I prefer map here largely because it's concise; if the mapping wasn't a preexisting function though, I'd use the genexpr to avoid map+lambda, which is the worst).