【Tips】Python 针对函数返回值实现列表推导

背景

设想这样一个场景，我们定义函数foo来生成一个大于0的随机code。但是函数执行可能会出现异常，对于异常情况我们指定函数返回-1。

在批量处理的情况下，我们想要得到所有成功执行的foo函数的返回值集合，而忽略返回-1的情况。基本的逻辑如下：

result = []

for item in item_lst:
    code = foo(item)
    if code != -1:
        result.append(code)

使用列表推导的解决方案

如果想用列表推导的方式实现，有下面几种解决方案：

Solution 1:

result = [foo(item) for item in item_lst]
result = [code for code in result if code != -1]

先用列表推导获取到了函数的返回值，然后再执行一次筛选，需要执行两次列表推导。

Solution 2:

result = [foo(item) for item in item_lst if foo(item) != -1]

执行一次列表推导，但是很明显，foo函数执行了两次。如果foo函数比较耗时，或者多次执行会产生副作用的话都会有问题。

Solution 3:

result = [code for code in (foo(item) for item in item_lst) if code != -1]

使用 生成器表达式 来处理中间结果，不需要临时列表，并且是基于迭代器的，优于上面两种方法。

参考

1. How to execute a filter by function return value in list comprehension?