Python含if-else的列表推导式未按预期工作问题排查

问题重现 #

Customer类定义 #

class Customer:
    c_id = None
    
    def __init__(self, cid):
        self.c_id = cid

创建对象与测试列表 #

a = Customer(1)
b = Customer(2)

lst = [a,b]

两个版本的函数 #

# Ver 1（报错版本）
def get_ids_v1(obj):
    return [obj.c_id if not isinstance(obj, list) else cust.c_id for cust in obj]

# Ver 2（正常版本）
def get_ids_v2(obj):
    if isinstance(obj, list):
        return [cust.c_id for cust in obj]
    else:
        return [obj.c_id]

测试结果 #

Ver 2运行正常：

get_ids_v2(a)
# Out: [1]

get_ids_v2(lst)
# Out: [1, 2]

Ver 1传入列表正常，但传入Customer对象时报错：

get_ids_v1(lst)
# Out: [1, 2]

get_ids_v1(a)
# Out: TypeError: 'Customer' object is not iterable

对比字典实现（正常运行） #

def get_id(obj):
    return [obj["id"] if not isinstance(obj, list) else cust["id"] for cust in obj]

c = {"id":3}
d = [{"id":4}, {"id":5}]

get_id(c)
# Out: [3]

get_id(d)
# Out: [4, 5]

报错原因分析 #

Ver 1的列表推导式逻辑存在本质问题：
列表推导式的结构是[表达式 for 变量 in 可迭代对象]，循环部分for cust in obj会优先于条件判断执行，并非按预期“条件成立就跳过循环”。

当传入Customer对象a时：

1. 首先执行for cust in obj，尝试遍历a这个Customer实例，但你的Customer类未实现迭代器协议，直接抛出TypeError: 'Customer' object is not iterable。
2. 此时条件判断if not isinstance(obj, list)的逻辑完全没机会执行，因为循环遍历已经先报错。

而字典实现能正常运行的原因是：
字典本身是可迭代对象（默认遍历键），当传入字典c时：

1. for cust in obj会遍历字典的键，循环一次后cust的值为"id"。
2. 条件判断if not isinstance(obj, list)成立，执行obj["id"]得到3，生成列表[3]，else分支的cust["id"]不会被执行，因此无报错。

修正方案 #

如果想用列表推导式实现类似逻辑，可先通过三元表达式统一迭代对象，确保其可迭代：

def get_ids_v1_fixed(obj):
    return [item.c_id for item in (obj if isinstance(obj, list) else [obj])]

无论输入是单个对象还是列表，都会先转换成合法的可迭代列表，再遍历提取c_id。

内容的提问来源于stack exchange，提问作者Prathamesh Pavnoji