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 循环。代码的简洁性只有在易于阅读时才有价值。