Python list comprehension: test function return

Question

Is there a way to test the return of a function in a list (or dict) comprehension? I'd like to avoid writing that:

lst = []
for x in range(10):
  bar = foo(x)
  if bar:
    lst.append(bar)

and use a list comprehension instead. Obviously, I don't want to write:

[foo(x) for x in range(10) if foo(x)]

so?

[foo(x) for x in range(10) if ??? ]

Answer 1

How about

filter(None, map(foo, range(10)))

If you don't want to keep the intermediate list, replace map() with itertools.imap(). And with itertools.ifilter(), the whole thing could be turned into a generator.

itertools.ifilter(None, itertools.imap(foo, range(10)))

Answer 2

Just make a generator to compute the values and build the filtered list from the generator afterwards.

Example:

# create generator via generator expression
results = (foo(x) for x in xrange(10))
# build result list, not including falsy values
filtered_results = [i for i in results if i]

For reference:

• https://docs.python.org/3/glossary.html#term-generator
• https://wiki.python.org/moin/Generators
• https://www.python.org/dev/peps/pep-0289/