Search code examples
python-3.xif-statementlist-comprehension

Python list comprehension ignore None results?


I have the following toy example function and list comprehension:

def foo(lst):
    if lst:
        return lst[0] + lst[1]

[foo(l) for l in [[], [1,2], [1,4]]]

The result is:

[None, 3, 4]

How can I avoid the Nones, I would like to avoid calling if foo(l) is not None inside the list comp. Please advise.


Solution

  • If you want to avoid calling the function more than once, you can make a generator that yields based on the result of the function. It's a little more code, but avoids making a list with a bunch of None values which have to be filtered later, and also avoids calling the function twice in the list comprehension:

    def expensive_function(lst):
        # sometimes returns None...hard to know when without calling
        if lst:
            return lst[0] + lst[1]
    
    def gen_results(l):
        for a in l:
            res = expensive_function(a)
            if res:
                yield res
        
    inp = [[], [1,2], [1,4]]
    
    list(gen_results(inp))
    # [3, 5]
    

    Also, since generators are lazy, you don't need to make a list if you don't need a list.