How to set local variable in list comprehension?

Question:

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] 
Asked By: Hao Tan

||

Answers:

You can avoid re-calculation by using python built-in filter:

list(filter(lambda t: t is not None, map(map_to_obj, v_list)))
Answered By: behzad.nouri

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]
Answered By: Lying Dog

List comprehensions are fine for the simple cases, but sometimes a plain old for loop is the simplest solution:

other_list = []
for v in v_list:
    obj = map_to_obj(v)
    if obj:
        other_list.append(obj)

Now if you really want a list comp and dont want to build an tmp list, you can use the iterator versions of filter and map:

import itertools as it
result = list(it.ifilter(None, it.imap(map_to_obj, v_list)))

or more simply :

import itertools as it
result = filter(None, it.imap(map_to_obj, v_list)))

The iterator versions don’t build a temporary list, they use lazy evaluation.

Answered By: bruno desthuilliers

I have figured out a way of using reduce:

def map_and_append(lst, v):
    mapped = map_to_obj(v)
    if mapped is not None:
        lst.append(mapped)
    return lst

reduce(map_and_append, v_list, [])

How about the performance of this?

Answered By: Hao Tan

A local variable can be set within a comprehension by cheating a bit and using an extra ‘for’ which “iterates” through a 1-element tuple containing the desired value for the local variable. Here’s a solution to the OP’s problem using this approach:

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

Here, o is the local variable being set equal to map_to_obj(v) for each v.

In my tests this is slightly faster than Lying Dog’s nested generator expression (and also faster than the OP’s double-call to map_to_obj(v), which, surprisingly, can be faster than the nested generator expression if the map_to_obj function isn’t too slow).

Answered By: Ovaflo

A variable assignment is just a singular binding:

[x   for v in l   for x in [v]]

This is a more general answer and also closer to what you proposed.
So for your problem you can write:

[x   for v in v_list   for x in [map_to_obj(v)]   if x]
Answered By: Vincent Goossens

Starting in Python 3.8, and the introduction of assignment expressions (PEP 572) (:= operator), it’s possible to use a local variable within a list comprehension in order to avoid calling the same function twice.

In our case, we can name the evaluation of map_to_obj(v) as a variable o while using the result of the expression to filter the list; and thus use o as the mapped value:

[o for v in [v1, v2, v3, v4] if (o := map_to_obj(v))]
Answered By: Xavier Guihot

The best way I have found to do this, which is cryptic but succinct is like this:

[f(y) for x in some_list()
      for y in [some_item(x)]  # wastefully making a singleton list
      if some_condition(y)]
Answered By: Jim Newton
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.