ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Python迭代器与生成器:从for循环本质到内存优化实践

Python迭代器与生成器:从for循环本质到内存优化实践 1. 从一道面试题说起for循环到底在做什么1.1 面试时我期待的答案语法糖拆解如果有人问我Python里最被低估的语法糖是什么我会毫不犹豫地说for循环。很多人写了三年Python每天和for循环打交道却不知道这行简单的代码背后藏着一整套迭代器协议。先看一个最普通的例子lst [1, 2, 3] for i in lst: print(i)如果面试官问你“Python解释器是怎么执行这段代码的”你会怎么回答如果只说“把列表里的元素一个一个取出来”那只能算直觉式理解谈不上机制。真正的答案是Python解释器拿到这个可迭代对象后会先调用iter(lst)获得一个迭代器对象然后在一个隐式的循环里反复调用next()取得下一个值直到捕获到StopIteration异常才结束整个循环。也就是说for其实是一个语法糖真正干活的是迭代器协议。这个协议并不是什么高深理论它就在你每天写的代码里只是你很少直视它。理解它不仅能让你看懂for循环的本质还能帮你写出更优雅、更省内存的代码也能让你在面对“为什么迭代器只能用一次”这类问题时不再迷茫。1.2 可迭代对象和迭代器别傻傻分不清很多初学者会把“可迭代对象”和“迭代器”当成同一个东西其实两者有明确区别。可迭代对象实现了__iter__方法调用后会返回一个迭代器。常见的列表、元组、字典、集合、字符串都是可迭代对象。迭代器实现了__iter__和__next__两个方法。其中__iter__返回自身__next__返回下一个值没有值了就抛出StopIteration。简单来说可迭代对象是“原料”迭代器是“加工后的游标”。你可以反复从同一条列表创建多个迭代器但每个迭代器只能一路走到黑。用代码验证一下lst [1, 2, 3] print(hasattr(lst, __next__)) # False列表不是迭代器 it iter(lst) print(hasattr(it, __next__)) # Trueiter() 返回迭代器我把两者的特性放在一张表里方便你对照特性可迭代对象迭代器实现__iter__是是实现__next__否是能直接用于for循环是是遍历后会耗尽否是示例list, str, dict, setiter(list), 生成器这张表解决了我当年的一大困惑为什么列表可以被for循环遍历多次而文件对象读完之后就没了因为列表是可迭代对象每次for循环都会重新创建一个新迭代器文件对象本身就是一个迭代器遍历完就走到头了。2. 迭代器协议__iter__和__next__的秘密约定2.1 手写一个迭代器协议最小实现迭代器协议只有两个方法写起来非常轻量。假如我想实现一个倒计时迭代器从指定数字递减到1class CountDown: def __init__(self, start): self.current start def __iter__(self): return self def __next__(self): if self.current 0: raise StopIteration val self.current self.current - 1 return val for num in CountDown(3): print(num) # 输出3 2 1注意这里的__iter__返回了self也就是说这个类本身就是迭代器。这是一种常见写法但如果你设计的是一个容器类比如一个自定义的集合那么__iter__应该返回一个新的迭代器对象而不是返回self。原因很简单迭代器是有状态的如果多个循环共享同一个状态第二个for循环就会直接结束。你可能会问for循环怎么知道什么时候停止答案是__next__抛出StopIteration异常。你没听错Python是用“异常”来结束循环的这在直觉上有点反常但却是整个协议最优雅的设计。2.2 为什么结束信号是StopIteration而不是None刚开始接触迭代器的人很容易想到一个问题为什么不能用return None来表示“没有更多元素了”假设你在遍历一批传感器数据其中某个值就是None表示“这次没有读数”。如果迭代器用None作为结束信号那么你遇到一个合法的None数据时循环就被误终止了这显然不能接受。StopIteration是一个独立于数据空间的异常信号它不占任何值的坑天然不会被数据干扰。另外异常还能携带额外信息。在Python 3.7中StopIteration可以带一个value属性用于在子生成器结束返回值时传递数据。这就是为什么yield from能够把子生成器的返回值优雅地传给外部。这里顺便说一句StopIteration也是协程间通信的一部分。如果你用next(gen)手动推进生成器看到StopIteration弹出不要慌这说明生成器已经运行完了。2.3 用while循环复刻一个for如果我们不依赖for语法手动用while来模拟它逻辑会清晰很多。下面这个my_for函数接受一个可迭代对象和一个回调函数def my_for(iterable, func): it iter(iterable) while True: try: value next(it) except StopIteration: break func(value) my_for([1, 2, 3], print) # 输出1 2 3代码一共只有几行却完整复现了for循环的核心逻辑先取迭代器再循环调用next()遇到StopIteration就跳出。理解了这段代码你就不会再觉得for循环是什么黑魔法了。3. 生成器迭代器最优雅的工厂3.1 yield是如何改写函数行为的虽然手写迭代器类不难但每次都要定义两个方法写多了也嫌烦。Python提供了一个更简洁的工具生成器。一个函数里只要出现yield关键字它就变成了生成器函数。调用生成器函数时函数体不会立刻执行而是返回一个生成器对象这个对象天然就是迭代器。def count_down(start): while start 0: yield start start - 1 for num in count_down(3): print(num) # 输出3 2 1yield的作用可以理解为“暂停并返回”。当next()执行到yield时函数会把值返回给调用方同时记住当前位置下一次next()时函数从上次暂停的位置继续运行直到再次遇到yield或函数结束。函数结束时Python会自动抛出StopIteration。这种机制让生成器成为编写“惰性序列”的完美工具。你不需要一次性把所有值算出来而是按需生产用多少算多少。3.2 生成器表达式与列表推导式的内存对决生成器表达式是生成器的一种极简写法和列表推导式只差一个括号。对比一下list_comp [x * x for x in range(1000000)] # 立刻生成100万个元素 gen_expr (x * x for x in range(1000000)) # 生成器对象几乎不占内存列表推导式会主动创建一个包含100万个元素的列表如果只是遍历一次这100万个元素会长时间占用内存。生成器表达式则是一个惰性对象每次迭代时才计算下一个值内存占用几乎可以忽略不计。我遇到过不少同事不管数据规模多大都习惯用列表推导式结果偶尔爆内存。我的建议很简单如果你只是“用一次”而且数据量较大优先考虑生成器表达式如果你需要反复遍历、随机访问或者数据量很小再考虑列表或元组。3.3 send与yield from生成器的高级玩法生成器除了“产出值”还能“接收值”这是通过send()实现的。看这个例子def echo(): while True: value yield print(f收到: {value}) gen echo() next(gen) # 启动生成器执行到 yield gen.send(hello) # 向生成器发送值 # 收到: hellosend()能把数据“喂”回生成器让生成器根据外部输入改变内部逻辑。这种用法在协程和状态机中很常见虽然平时写业务代码不一定用得上但理解了它会让你对“生成器是双向通道”有更深的认识。另一个高频用法是yield from它可以委托一个子生成器。比如我们要把多个可迭代对象串起来def chain(*iterables): for it in iterables: yield from it list(chain(ab, [1, 2])) # [a, b, 1, 2]yield from不仅会转发子生成器的所有产出值还能处理异常的传递代码比手动嵌套for循环简洁得多。4. 真实项目里的迭代器惰性求值与手写数据流4.1 文件对象就是最好的迭代器例子可能你每天都在用文件迭代器却没有意识到它有多优秀。比如读取一个超大日志文件with open(huge.log) as f: for line in f: process(line)这里的f就是文件迭代器。for循环逐行读取文件内容每次只在内存中保留一行。哪怕文件有几个GB程序也能稳定运行。如果你改成f.readlines()那就要一次性把整个文件载入内存分分钟把内存吃光。迭代器的惰性求值能力让“流式处理”成为可能。不仅是文件数据库游标、网络响应流、消息队列消息这些本质上都是“流”都适合用迭代器来封装。4.2 迭代器的一次性陷阱与三种解法迭代器最容易被误用的特性就是“一次性”。看这段代码it iter([1, 2, 3]) print(list(it)) # [1, 2, 3] print(list(it)) # []第二次list(it)拿到的是空列表因为迭代器已经走到底了。遇到需要重复遍历的场景通常有三种解法数据量小直接转成列表或元组缓存之后随便遍历几次都行。数据量大用itertools.tee把迭代器复制成多个独立分支但要注意tee内部会缓存数据源迭代器消耗越大缓存越大。设计成工厂函数每次需要迭代时就重新创建迭代器。比如写一个函数返回生成器而不是返回生成器对象本身。我个人的原则是如果这个数据会被遍历超过一次我就倾向于把它转成列表除非它大到内存装不下。用“能不能重复遍历”这个标准来设计接口能避免很多隐藏Bug。4.3 自定义可回退迭代器一个实战案例有一次我做数据同步从数据库A读取一批记录清洗后写入数据库B。因为数据量很大我不能一次性全部读出但写入时如果某条记录有问题需要回退到上一条重新处理。这时候需要一个支持“回退”的迭代器。我写了一个简单的RewindableIteratorclass RewindableIterator: def __init__(self, iterable): self._it iter(iterable) self._last None self._rewound False def __iter__(self): return self def __next__(self): if self._rewound: self._rewound False return self._last self._last next(self._it) return self._last def rewind(self): self._rewound True使用方法很直接data [1, 2, 3, 4] it RewindableIterator(data) print(next(it)) # 1 print(next(it)) # 2 it.rewind() # 回退一步 print(next(it)) # 2 再次读取这个例子给了我一个启发迭代器协议并不限制你只能“往前”。你可以通过包装、缓存等方式扩展它的行为。在真实项目中迭代器往往不是单独存在的而是像管道一样一层包一层def read_sensor_data(): raw [21.5, 22.0, 999.9, 21.8, 22.1] for value in raw: yield value def filter_invalid(data_iter, threshold100): for value in data_iter: if value threshold: yield value for valid_value in filter_invalid(read_sensor_data()): print(valid_value) # 21.5 22.0 21.8 22.1这个“生成器管道”模式非常实用。每一层只做一件事数据按需流动既节省内存又方便单独测试。5. 迭代器与for循环的底层细节CPython的解释器行为5.1 用dis模块看for循环的字节码说到“揭秘”只看源码还是不过瘾我建议你看看字节码。Python的dis模块可以把函数编译后的字节码打印出来import dis def demo(lst): for i in lst: print(i) dis.dis(demo)输出简化后是4 0 LOAD_FAST 0 (lst) 2 GET_ITER 4 FOR_ITER 10 (to 16) 6 STORE_FAST 1 (i) 8 LOAD_GLOBAL 0 (print) 10 LOAD_FAST 1 (i) 12 CALL_FUNCTION 1 14 POP_TOP 16 JUMP_ABSOLUTE 4 18 LOAD_CONST 0 (None) 20 RETURN_VALUE注意GET_ITER和FOR_ITER这两条指令GET_ITER会对栈顶对象调用iter()得到迭代器FOR_ITER则负责调用next()如果遇到StopIteration就跳转到循环结束位置。整个循环不需要索引变量也不需要判断长度一切由迭代器协议驱动。理解了这一点你就明白为什么Python的for循环能遍历文件、字典、生成器、集合等等五花八门的类型——因为它们都遵守同一个协议。这也是“鸭子类型”的体现只要你有__iter__和__next__我就拿你当迭代器用。5.2 for循环和while循环谁更快很多人想当然地以为for循环速度不如while因为每次都要经历迭代器协议。但实际上在CPython中for循环通常比while循环更快。原因在于for循环的迭代逻辑大部分是在C语言层面实现的而while循环里的条件判断和变量更新都是在Python字节码层面执行的解释器需要翻译更多的字节码指令。当然这个差异在绝大多数场景下可以忽略不计。真正影响性能的关键还是循环体内部做的事情。如果循环体里有IO操作或复杂计算循环方式的选择反而不重要了。5.3 自定义迭代器时最容易踩的坑写自定义迭代器时有几个坑经常会让人卡半天。第一忘了实现__iter__。如果你只写了__next__然后直接拿去for循环会得到TypeError: X object is not iterable。因为Python发现对象没有__iter__就会退而求其次检查序列协议__getitem__但普通类通常也没实现。解决办法很简单把__iter__补上返回self或一个新的迭代器对象。第二迭代器被多个循环共享。如果__iter__返回self两个for循环用同一个对象时第二个循环会直接结束。我之前写一个容器类时就栽过这个跟头。正确的做法是让容器类的__iter__每次都返回一个全新迭代器。第三在遍历列表时删除元素。for i in lst: lst.remove(i)这种代码极其危险因为列表内部索引在变化你会莫名跳过后面的元素。正确做法是用列表推导式或filter生成新列表lst [x for x in lst if keep(x)]。遇到迭代器相关的诡异问题我习惯先写一个带调试输出的包装器把每次next()的值打出来很快就能定位问题出在哪个环节。6. 迭代器进阶itertools工具箱与常见陷阱6.1 五个高频itertools函数与示例itertools是标准库里最被低估的模块之一。它不是一堆花哨的装饰而是实实在在的“迭代器乐高”。我平时用得最频繁的几个函数chain把多个可迭代对象串成一个。示例import itertools for x in itertools.chain([1, 2], [3, 4]): print(x) # 1 2 3 4islice对迭代器做切片不会一次性加载全部数据。示例只读取文件前10行。with open(huge.log) as f: for line in itertools.islice(f, 10): print(line.rstrip())注意islice不支持负数步长这和列表切片不同。count无限计数器。通常搭配islice使用for num in itertools.islice(itertools.count(0, 2), 5): print(num) # 0 2 4 6 8不要直接for i in itertools.count()因为永远不会结束。groupby把连续相同的key分组。注意它只对连续元素有效如果数据乱序需要先排序。data [(a, 1), (a, 2), (b, 3)] for key, group in itertools.groupby(data, keylambda x: x[0]): print(key, list(group)) # a [(a, 1), (a, 2)] # b [(b, 3)]product笛卡尔积用来替代多层嵌套循环。for x, y in itertools.product(range(3), range(3)): print(x, y)多层嵌套循环会让人头皮发麻product能让代码扁平化而且配合break跳出多层循环时也更方便。6.2 迭代器使用的三大“坑中之坑”除了前面提到的“一次性”陷阱还有几个坑值得单独拿出来说。第一个坑tee复制后源迭代器被提前用完。看这段代码it iter(range(5)) a, b itertools.tee(it) print(list(it)) # 源迭代器已经被消耗掉 print(list(a)) # 空 print(list(b)) # 空这是因为tee是在源迭代器的基础上做缓存的如果一开始就把源迭代器消耗完复制出来的分支自然也没数据。正确做法是先tee再消耗源迭代器。第二个坑嵌套迭代器共享生成器状态。比如gen (i for i in range(5)) iter_list [iter(gen), iter(gen)] print(next(iter_list[0])) # 0 print(next(iter_list[1])) # 1两个iter对象实际上指向同一个生成器因为它们共享同一个内部状态。这提醒我们在用迭代器构造新迭代器时一定要确认底层数据源是否可以重复迭代。第三个坑生成器内部资源没有释放。生成器如果在yield之后被异常打断或直接close()如果内部没有finally块文件句柄等资源可能会泄漏。建议在生成器内部用with管理资源这是最稳妥的方式。6.3 最后的实用建议优先用生成器实现__iter__如果你实在不知道什么时候该用手写迭代器类什么时候该用生成器我给你的建议是在实现__iter__方法时优先用生成器。比如你想定义一个支持遍历的集合类class MyCollection: def __init__(self, data): self.data data def __iter__(self): for item in self.data: yield item这样__iter__返回的是一个生成器对象天然就是迭代器完全满足迭代器协议而且代码比手写一个反复next()的类要短得多。只有在需要维护复杂状态、支持回退、或者需要实现双向通信时才值得去手写一个完整的迭代器类。我自己的习惯是能用生成器解决的绝不去手写类。生成器让迭代器的创建成本降到了“写一个普通函数”的级别这也是Python最让我上瘾的地方。这个内容讲到这里已经把我对于迭代器协议、for循环、生成器和itertools的理解都倒出来了。以后你写for x in data的时候不妨多想一想这个data是迭代器还是可迭代对象它会不会在循环中被耗尽如果这些问题你能迅速回答那今天这篇文章就没白看。
返回列表