Python 中的 List、Dict、Set 推导式:简洁又易读

1. 概述

推导式是 Python 的特色语法。它允许你在一行代码中从一个可迭代对象创建 list、dict 或 set。写得好的话,代码比 for 循环更短更清楚。滥用的话,就会变成一长串难以理解的代码。

本文通过具体例子介绍 list、dict、set 推导式和生成器表达式。读完本文,你将了解何时应该使用推导式、何时应该回到 for 循环,以及如何为数据转换逻辑编写测试。所有示例都在 Python 3.12 上运行。

2. List 推导式

2.1. 基本语法

假设你有一个订单列表,需要获取每个订单的总金额。用传统 for 循环的写法如下:

from dataclasses import dataclass

@dataclass
class Order:
    id: int
    customer: str
    total: float
    status: str

orders = [
    Order(1, "An", 250_000, "PAID"),
    Order(2, "Bình", 90_000, "CANCELLED"),
    Order(3, "An", 1_200_000, "PAID"),
    Order(4, "Chi", 450_000, "PENDING"),
]

totals = []
for order in orders:
    totals.append(order.total)
print(totals)

用 list 推导式改写:

totals = [order.total for order in orders]
print(totals)
[250000, 90000, 1200000, 450000]

结构是 [表达式 for 元素 in 可迭代对象]。从左到右读:「对 orders 中的每个 order,取其 order.total」。你不需要声明空列表,也不需要调用 append。「我在创建一个新列表」这个意图从 [ 就已经很明显了。

2.2. 用 if 过滤元素

如果只想获取已支付的订单,在末尾添加 if 子句:

paid_ids = [order.id for order in orders if order.status == "PAID"]
print(paid_ids)
[1, 3]

if 位于末尾,起到过滤器的作用。不满足条件的元素会被过滤掉,所以结果列表可能比原列表更短。

2.3. 在前面放 if-else 表达式

当你想根据条件转换每个元素而不是过滤时,把 if-else 放在表达式部分:

labels = ["大" if order.total >= 500_000 else "小" for order in orders]
print(labels)
['小', '小', '大', '小']

这是很多人容易混淆的地方。你只需记住两个规则:

  • if 在 for 后面用来过滤。这个子句没有 else。
  • a if cond else b 在 for 前面是条件表达式。这个子句必须有 else,并且保持元素数量不变。

写成 [x for x in data if cond else y] 会导致 SyntaxError。

3. Dict 推导式

Dict 推导式使用 {} 和 key: value 对。例如,你需要能根据 id 快速查询订单:

orders_by_id = {order.id: order for order in orders}
print(orders_by_id[3])
Order(id=3, customer='An', total=1200000, status='PAID')

你也可以转换一个已有的 dict。例如,给价格表加上 10% 的税:

prices = {"咖啡": 30_000, "奶茶": 45_000, "面包": 20_000}
prices_with_vat = {name: round(price * 1.1) for name, price in prices.items()}
print(prices_with_vat)
{'咖啡': 33000, '奶茶': 49500, '面包': 22000}

prices.items() 返回 (key, value) 对,你可以在 for 循环中直接解包。

一个有用的技巧是反转 dict:

status_codes = {"PAID": 1, "PENDING": 2, "CANCELLED": 3}
code_to_status = {code: status for status, code in status_codes.items()}
print(code_to_status[2])
PENDING

注意: 如果两个元素生成相同的 key,后面的元素会静默覆盖前面的元素。例如 {order.customer: order.total for order in orders} 只会保留「An」最后的订单(1,200,000),250,000 的订单会无声地消失。需要分组时,请使用 defaultdict 配合 for 循环(见第 7 节)。

4. Set 推导式

Set 推导式也使用 {},但只有一个表达式,没有冒号 :。结果会自动去重。例如,你想知道哪些客户已经下过订单:

customers = {order.customer for order in orders}
print(sorted(customers))
['An', 'Bình', 'Chi']

每个客户只出现一次,尽管「An」有两个订单。Set 中的顺序不被保证,所以上例在打印时使用了 sorted。

当需要快速检查成员时,Set 推导式非常方便:

banned_emails = {"spam@x.com", "Fake@Y.com"}
normalized = {email.strip().lower() for email in banned_emails}
print("fake@y.com" in normalized)
True

注意: 空的 {} 是dict,不是 set。要创建空 set,需要写 set()。

5. 多个 for 和嵌套推导式

你可以在一个推导式中写多个 for。阅读顺序完全相同于从上到下写嵌套 for 循环的顺序:

carts = {
    "An": ["咖啡", "面包"],
    "Chi": ["奶茶"],
}

items = [(customer, item) for customer, products in carts.items() for item in products]
print(items)
[('An', '咖啡'), ('An', '面包'), ('Chi', '奶茶')]

上面的推导式等价于下面的代码:

items = []
for customer, products in carts.items():
    for item in products:
        items.append((customer, item))

最常见的情况是展平(flatten)一个嵌套列表:

matrix = [[1, 2, 3], [4, 5, 6]]
flat = [value for row in matrix for value in row]
print(flat)
[1, 2, 3, 4, 5, 6]

记忆方法:最外层的 for 循环放在前面。如果你需要读第二遍才能理解顺序,这就是应该分解成普通 for 循环的信号。

6. 生成器表达式:当不需要整个列表时

如果把 [] 换成 (),你就得到了一个生成器表达式。Python 不会在内存中创建列表。值会按需逐个生成。

revenue = sum(order.total for order in orders if order.status == "PAID")
print(revenue)
1450000

当生成器是函数的唯一参数时,你可以省略一对括号。这种写法适用于 sum、any、all、max、min:

has_big_order = any(order.total > 1_000_000 for order in orders)
all_positive = all(order.total > 0 for order in orders)
print(has_big_order, all_positive)
True True

any 一旦遇到第一个满足条件的元素就停止。因此对于大数据集,生成器既节省内存,也比 any([...]) 快。

注意: 生成器只能被遍历一次。第二次遍历不会返回任何东西:

paid = (o for o in orders if o.status == "PAID")
print(len(list(paid)))
print(len(list(paid)))
2
0

如果你需要多次使用结果,请使用 list 推导式。

7. 保持推导式易读

7.1. 为复杂逻辑命名函数

长条件会使推导式难以阅读:

vip = [o.customer for o in orders if o.status == "PAID" and o.total >= 1_000_000 and not o.customer.startswith("test")]

你可以把条件提取到一个有清晰名字的函数中:

def is_vip_order(order: Order) -> bool:
    return (
        order.status == "PAID"
        and order.total >= 1_000_000
        and not order.customer.startswith("test")
    )

vip = [order.customer for order in orders if is_vip_order(order)]
print(vip)
['An']

现在推导式读起来像一句完整的话,而业务逻辑也有独立的地方来测试。

7.2. 合理换行

推导式不一定要写在一行里。在括号内,你可以自由换行:

summary = {
    order.id: order.total
    for order in orders
    if order.status != "CANCELLED"
}
print(summary)
{1: 250000, 3: 1200000, 4: 450000}

每行一个角色:表达式、源、过滤器。

7.3. 用 Walrus 操作符避免重复计算

有时你既需要过滤,又需要使用同一个计算的结果。操作符 :=(Python 3.8+)可以让你只计算一次:

def discount_for(order: Order) -> float:
    return order.total * 0.05 if order.total >= 400_000 else 0

discounts = {
    order.id: d
    for order in orders
    if (d := discount_for(order)) > 0
}
print(discounts)
{3: 60000.0, 4: 22500.0}

要节制地使用 walrus 操作符。变量 d 会「泄漏」到推导式外的作用域。而且,很多读者还不熟悉这个语法。

8. 测试数据转换逻辑

推导式通常位于数据转换函数中,所以你应该像测试其他逻辑一样测试它们。首先,安装 pytest:

pip install pytest

假设你有一个 reports.py 模块:

# reports.py
from dataclasses import dataclass

@dataclass
class Order:
    id: int
    customer: str
    total: float
    status: str

def revenue_by_customer(orders: list[Order]) -> dict[str, float]:
    customers = {o.customer for o in orders if o.status == "PAID"}
    return {
        c: sum(o.total for o in orders if o.customer == c and o.status == "PAID")
        for c in customers
    }

以及对应的测试文件:

# test_reports.py
from reports import Order, revenue_by_customer

def test_revenue_only_counts_paid_orders():
    orders = [
        Order(1, "An", 250_000, "PAID"),
        Order(2, "An", 1_200_000, "PAID"),
        Order(3, "Bình", 90_000, "CANCELLED"),
        Order(4, "Chi", 450_000, "PAID"),
    ]

    result = revenue_by_customer(orders)

    assert result == {"An": 1_450_000, "Chi": 450_000}
    assert "Bình" not in result

def test_empty_orders_returns_empty_dict():
    assert revenue_by_customer([]) == {}
pytest -q
..                                                                       [100%]
2 passed in 0.01s

第一个测试断言两点:取消的订单不计入,只有取消订单的客户不出现在结果中。第二个测试检查边界情况(空列表)。

上面的函数是正确的,但时间复杂度是 O(n × 客户数),因为要为每个客户遍历全部 orders。对于大数据集,结合 defaultdict 的 for 循环既更快又更易读:

from collections import defaultdict

def revenue_by_customer(orders: list[Order]) -> dict[str, float]:
    totals: dict[str, float] = defaultdict(float)
    for order in orders:
        if order.status == "PAID":
            totals[order.customer] += order.total
    return dict(totals)

两个旧测试仍然通过,无需修改。这就是有测试的好处:你可以放心地在推导式和 for 循环之间重构。

9. 何时不应该使用推导式

  • 需要副作用时。 不要写 [print(o) for o in orders] 或 [db.save(o) for o in orders]。这样会产生一个无意义的 None 列表,并误导读者认为你需要结果。请使用 for 循环。
  • 嵌套超过两层或有多个交错的 if。一行中三个 for 循环几乎总是比分开写的版本难以阅读。
  • 需要分组或累积,如求和或计数。请使用 defaultdict、Counter 或 for 循环。
  • 需要处理每个元素的异常时。 推导式不能包含 try/except。你应该把它提取到单独的函数中然后在推导式中调用,或者使用 for 循环。
  • 数据非常大且只遍历一次时。 此时应该用生成器表达式替代 list。

注意: 在 Python 3 中,推导式的循环变量不会泄漏到外部作用域。执行完 [x for x in range(3)] 后,名字 x 不存在于外部作用域,这与普通 for 循环不同。唯一的例外是用 := 赋值的变量,如第 7.3 节所述。

10. 总结

推导式帮助你用简洁的方式表达「从旧集合创建新集合」这一意图:

  • [...] 用于 list,{k: v ...} 用于 dict,{...} 用于 set,(...) 用于生成器。
  • if 在 for 后面用来过滤,a if cond else b 在 for 前面用来转换。
  • 把复杂的条件提取到命名清晰的函数中,表达式长时换行。
  • 用生成器表达式配合 sum、any、all 来节省内存。
  • 当需要处理副作用、分组、异常,或推导式开始难以阅读时,回到 for 循环。

最简单的规则:如果同事需要读两遍你的推导式才能理解,请改写成 for 循环。代码的简洁性只有在易于阅读时才有价值。

pythoncomprehensionclean codegeneratorpytest