Python 迭代器与生成器

处理大数据时,最怕一次性把所有数据装进内存。迭代器与生成器就是 Python 的"懒加载"方案:按需产生数据,边算边用。理解它们,你的代码会更快、更省内存,也更 Pythonic。

1. 可迭代对象与迭代器

列表、字符串、字典都能 for 遍历,因为它们实现了 __iter__;迭代器则是"每次吐出一个值"的对象,用 iter() 和 next() 驱动:

nums = [1, 2, 3]
it = iter(nums)          # 拿到迭代器
print(next(it))          # 1
print(next(it))          # 2
print(next(it))          # 3
# 再调用 next(it) 会抛 StopIteration

2. for 循环的本质

for 循环就是不断调用 next() 直到捕获 StopIteration 的语法糖:

it = iter("abc")
while True:
    try:
        ch = next(it)
        print(ch)
    except StopIteration:
        break

3. 生成器函数:yield 的魔法

含 yield 的函数调用后返回生成器:每次 next 执行到 yield 暂停,再次 next 继续往下走:

def countdown(n):
    while n > 0:
        yield n
        n -= 1

for x in countdown(3):
    print(x)   # 3 2 1

生成器最大的好处是惰性:不一次性生成全部数据,内存占用恒定。

4. 生成器表达式

语法与列表推导式一样,只是用圆括号,返回生成器而非列表:

squares = (x * x for x in range(10))
print(sum(squares))     # 285,边算边求和

# 对比:列表推导式会先建出整个列表
big = [x for x in range(1_000_000)]        # 占内存
big_gen = (x for x in range(1_000_000))    # 几乎不占

5. 实战:逐行处理大文件

文件对象本身就是迭代器,逐行读取是处理 GB 级日志的标准姿势:

def errors_in_log(path):
    with open(path, encoding="utf-8") as f:
        for line in f:            # 惰性逐行,内存恒定
            if "ERROR" in line:
                yield line.strip()

for err in errors_in_log("app.log"):
    print(err)

6. itertools:迭代器工具箱

标准库 itertools 提供了大量迭代器工具,常用三个:

from itertools import islice, chain, groupby

# 只取前 3 个
first3 = list(islice(range(100), 3))      # [0, 1, 2]

# 拼接多个迭代器
merged = list(chain("ab", "cd"))          # ['a', 'b', 'c', 'd']

# 按 key 分组(数据需先排序)
data = sorted(["a1", "a2", "b1"], key=lambda s: s[0])
for key, group in groupby(data, key=lambda s: s[0]):
    print(key, list(group))

💡 学习建议:把之前"先建列表再处理"的代码改成生成器版本,用 tracemalloc 对比内存占用,你会真正理解惰性求值的威力。