Python Comprehension 和 Generator 技巧

本文最后更新于 22 天前，如有失效请评论区留言。

写 Python 的时候，你有没有遇到过这种情况：明明一个简单的数据转换，却要写好几行 for loop + append？或者处理一个超大文件，内存直接爆了？今天就来聊聊 Python 里两个特别好用的语法糖：comprehension 和 generator expression。掌握它们，你的代码会变得更简洁、更高效、更 Pythonic。

一、List Comprehension：一行顶三行

List comprehension 是 Python 最经典的语法特性之一。它的核心思想就是：用一行表达式来创建一个新 list。

1.1 基本语法

最简单的形式长这样：

# 传统写法
squares = []
for x in range(10):
    squares.append(x ** 2)

# List comprehension
squares = [x ** 2 for x in range(10)]
# [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

看起来就是把 for loop 压缩成了一行，对吧？但它不只是”语法糖”那么简单——在 CPython 实现中，list comprehension 比等价的 for loop + append 快，因为它底层用了专门的字节码指令，省去了每次循环调用 append() 方法的开销。

1.2 带条件过滤

你可以在 comprehension 里加 if 来过滤元素：

# 只保留偶数的平方
even_squares = [x ** 2 for x in range(10) if x % 2 == 0]
# [0, 4, 16, 36, 64]

# 过滤出字符串列表中的非空项
words = ["hello", "", "world", "", "python"]
non_empty = [w for w in words if w]
# ['hello', 'world', 'python']

1.3 嵌套 comprehension

嵌套 comprehension 可以处理二维数据，但要注意可读性：

# 展平一个二维列表
matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flat = [num for row in matrix for num in row]
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

# 生成坐标对
coords = [(x, y) for x in range(3) for y in range(3)]
# [(0,0), (0,1), (0,2), (1,0), (1,1), ...]

💡 小贴士：嵌套 comprehension 的阅读顺序是从左到右、从外到内。如果嵌套超过两层，建议还是用普通的 for loop，不然自己过几天回来都看不懂。

1.4 带 if-else 的 comprehension

注意，if-else 和单纯的 if 位置不一样：

# if-else 放在前面（三元表达式）
labels = ["even" if x % 2 == 0 else "odd" for x in range(5)]
# ['even', 'odd', 'even', 'odd', 'even']

# 纯 if 过滤放在后面
evens = [x for x in range(10) if x % 2 == 0]

记住这个规则：if-else 是表达式的一部分，放前面；纯 if 过滤放后面。

二、Dict Comprehension 和 Set Comprehension

不只是 list，dict 和 set 也有自己的 comprehension：

# Dict comprehension：快速构建字典
word_lengths = {word: len(word) for word in ["hello", "world", "python"]}
# {'hello': 5, 'world': 5, 'python': 6}

# 翻转字典的 key-value
original = {"a": 1, "b": 2, "c": 3}
flipped = {v: k for k, v in original.items()}
# {1: 'a', 2: 'b', 3: 'c'}

# Set comprehension：自动去重
text = "hello world"
unique_chars = {ch for ch in text if ch != ' '}
# {'h', 'e', 'l', 'o', 'w', 'r', 'd'}

Dict comprehension 在做数据转换的时候特别好用，比如把一个 list of tuples 变成 dict，或者做 key 的映射。

三、Generator Expression：惰性求值的威力

Generator expression 看起来和 list comprehension 几乎一样，只不过用的是圆括号 () 而不是方括号 []：

# List comprehension - 立即生成所有元素
squares_list = [x ** 2 for x in range(1000000)]

# Generator expression - 惰性生成，按需计算
squares_gen = (x ** 2 for x in range(1000000))

区别在哪？List comprehension 会一次性把所有结果算出来、放进内存；而 generator expression 不会——它返回一个 generator 对象，每次你 next() 它的时候才计算下一个值。

3.1 为什么要用 Generator？

两个字：省内存。

假设你要处理一个 1GB 的日志文件，逐行统计某个关键词出现的次数：

# ❌ 错误做法：把所有行读进内存
with open("huge_log.txt") as f:
    lines = f.readlines()  # 内存爆炸！
    count = sum(1 for line in lines if "ERROR" in line)

# ✅ 正确做法：逐行读取 + generator
with open("huge_log.txt") as f:
    count = sum(1 for line in f if "ERROR" in line)  # 文件对象本身就是 iterable

上面的 sum(1 for line in f if "ERROR" in line) 里，1 for line in f if "ERROR" in line 就是一个 generator expression。它每次只产生一个 1，sum() 加完就丢掉，内存占用恒定。

3.2 Generator 的常见用法

Generator expression 经常和接受 iterable 的内置函数配合使用：

# sum / min / max / any / all 都很适合搭配 generator
total = sum(x ** 2 for x in range(1000))
max_val = max(len(word) for word in words)
has_empty = any(not line.strip() for line in lines)
all_positive = all(x > 0 for x in numbers)

注意这里连外层的括号都省了——当 generator expression 是函数的唯一参数时，可以省略那对圆括号。

四、Generator Function：更强大的 generator

除了 generator expression，Python 还有 generator function——用 yield 关键字定义的函数。它比 generator expression 更灵活，适合复杂的生成逻辑。

# 生成斐波那契数列
def fibonacci():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

# 取前 10 个斐波那契数
from itertools import islice
fib_10 = list(islice(fibonacci(), 10))
# [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

这个 fibonacci() 函数返回一个无限的 generator。它不会把所有数都算出来（那也不可能），而是每次 next() 的时候才算下一个。配合 islice 就可以按需取用。

4.1 yield vs return

return 是一次性返回然后函数结束；yield 是「暂停」函数，保留当前状态，下次调用的时候从暂停的地方继续。

# 读取大文件的 chunk
def read_in_chunks(file_path, chunk_size=8192):
    with open(file_path, 'rb') as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            yield chunk

# 使用
for chunk in read_in_chunks("big_file.bin"):
    process(chunk)  # 每次只处理一小块

五、实战技巧和常见坑

5.1 Comprehension 里不要写副作用

# ❌ 千万别这么干
[print(x) for x in range(5)]

# ✅ 如果需要副作用，用普通 for loop
for x in range(5):
    print(x)

Comprehension 的目的是创建新数据，不是执行副作用。用它来打印日志、写文件、修改外部变量都是 anti-pattern。

5.2 Generator 只能遍历一次

gen = (x ** 2 for x in range(5))

list(gen)  # [0, 1, 4, 9, 16]
list(gen)  # []  ← 空了！已经遍历完了

Generator 是一次性的，用完就没了。如果你需要多次遍历，要么转成 list，要么重新创建 generator。

5.3 Walrus Operator（:=）在 comprehension 中的应用

Python 3.8 引入的 walrus operator 可以在 comprehension 中避免重复计算：

# 假设有一个耗时计算函数
results = [
    y
    for x in data
    if (y := expensive_function(x)) > threshold
]

这样 expensive_function(x) 只调用一次，结果同时用于过滤和赋值。

六、什么时候用什么？

需要一个新 list → list comprehension
只需要遍历一次、数据量大 → generator expression
生成逻辑复杂 → generator function（yield）
需要构建 dict/set → dict/set comprehension
嵌套超过两层 → 老老实实写 for loop

一句话总结：comprehension 是用来创建集合的，generator 是用来惰性遍历的。分清「创建」和「遍历」这两个场景，就不会选错。

总结

今天我们从 list comprehension 讲到 generator expression，再到 generator function，覆盖了 Python 中最常用的「推导式」和「生成器」技巧。这些特性是 Python 写起来简洁高效的核心原因之一。下次你写 for loop + append 的时候，想想是不是可以用 comprehension 来简化；处理大数据集的时候，想想是不是该用 generator 来省内存。

代码写得简洁不是为了炫技，而是为了让意图更清晰。Happy coding! 🐍