How to set local variable in list comprehension?

python list-comprehension

Hao Tan · Oct 31, 2014 · Viewed 29.7k times · Source

I have a method that takes a list and returns an object:

# input a list, returns an object
def map_to_obj(lst):
    a_list = f(lst)
    return a_list[0] if a_list else None

I want to get a list that contains all the mapped elements that aren't None.

Like this:

v_list = [v1, v2, v3, v4]

[map_to_obj(v) for v in v_list if map_to_obj(v)]

But it doesn't seem good to call the map_to_obj method twice in the list comprehension.

Is there a way to have local variables in list comprehensions so that it can have better performance?

Or does the compiler optimize it automatically?

Here is what I want:

(sml like)
[let mapped = map_to_obj(v) in for v in v_list if mapped end]

Answer

Use nested list comprehension:

[x for x in [map_to_obj(v) for v in v_list] if x]

or better still, a list comprehension around a generator expression:

[x for x in (map_to_obj(v) for v in v_list) if x]