speedplane
speedplane

Reputation: 16121

Python List Comprehension: Using "if" Statement on Result of the Comprehension

Can you filter a list comprehension based on the result of the transformation in the comprehension?

For example, suppose you want to strip each string in a list, and remove strings that are just whitespace. I could easily do the following:

filter(None, [x.strip() for x in str_list])

But that iterates over the list twice. Alternatively you could do the following:

[x.strip() for x in str_list if x.strip()]

But that implementation performs strip twice. I could also use a generator:

for x in str_list:
   x = x.strip()
   if x:
      yield x

But now that's a bunch of lines of code. Is there any way of doing the above (1) only iterating once; (2) only performing the transformation once; and (3) in a single list comprehension? The example above is a toy example, but I'd like to do this with longer lists and non-trivial transforms.

Update: I'm using Python 2.7.X, and prefer answers in that, but if Python 3 has some new features that make this easy, I'd be happy to learn about them as well.

Upvotes: 8

Views: 769

Answers (2)

Mad Physicist
Mad Physicist

Reputation: 114230

Don't pass a list to filter, pass a generator expression, and it will only be iterated once:

filter(None, (x.strip() for x in str_list))

This is exactly the same idea as using a nested generator like

[y for y in (x.strip() for x in str_list) if y]

Both cases rely on the lazy evaluation of generators: each element of str_list will be processed exactly once, when the corresponding output element is created. No intermediate lists will be made.

The comprehension approach is nice for small one-shot transformations like these. Even the simple example here, which filters after the transformation, is pushing the limits of readability in my opinion. With any non-trivial sequence of transformations and filters, I would recommend using a for loop.

Upvotes: 7

Filipp
Filipp

Reputation: 2040

Why not nest the comprehensions?

result = (x for x in (y.strip() for y in str_list) if len(x))

The inner () makes a generator that is just the stripped versions of strings in str_list. The outer () makes a second generator that consumes the first and produces only the non-empty elements. You traverse the list only once and strip each string only once.

Upvotes: 2

Related Questions