057、迭代器协议与自定义迭代器iter、next与 itertools 混用上周帮同事排查一个数据管道的内存泄漏问题日志里看到StopIteration异常被吞掉导致生成器无限循环。翻代码发现他手写了一个迭代器类__next__方法里忘了处理边界条件——这种坑我踩过不止一次。今天干脆把迭代器协议掰开揉碎讲清楚顺便聊聊怎么跟itertools配合着玩。迭代器协议到底是个啥Python 里迭代器协议就两条规则对象实现__iter__返回迭代器自身__next__每次返回下一个元素没元素时抛StopIteration。别小看这两条很多新手写自定义迭代器时要么忘了__iter__返回 self要么__next__里用return None代替抛异常——后者会导致循环无法正常终止。看个反面教材我当年也这么写过classBadRange:def__init__(self,n):self.nn self.i0def__iter__(self):returnself# 这里没问题def__next__(self):ifself.iself.n:valself.i self.i1returnval# 别这样写返回 None 会让 for 循环继续跑# return NoneraiseStopIteration# 必须抛这个for循环底层就是不断调__next__直到捕获StopIteration。你返回None循环会认为None是有效值继续调下一次——死循环就来了。可迭代对象 vs 迭代器别搞混很多教程把这两个概念混着讲实际区别很关键。可迭代对象比如列表、字符串实现了__iter__返回一个迭代器迭代器同时实现了__iter__和__next__。列表本身不是迭代器你调iter([1,2,3])才拿到迭代器。为什么要有这个区分因为迭代器是一次性消耗品遍历完就废了。可迭代对象可以反复生成新的迭代器。看这个例子nums[1,2,3]it1iter(nums)it2iter(nums)# 每次调 iter() 得到新迭代器print(list(it1))# [1, 2, 3]print(list(it1))# [] it1 已经耗尽print(list(it2))# [1, 2, 3] it2 是新的如果你把列表本身当迭代器用比如直接给next()会报TypeError。这里踩过坑写代码时图省事直接把列表传给next()结果运行时炸了。自定义迭代器的正确姿势写一个能反复遍历的自定义迭代器需要把迭代器对象和可迭代对象分开。比如实现一个“斐波那契数列迭代器”classFibs:def__init__(self,max_count):self.max_countmax_countdef__iter__(self):returnFibIterator(self.max_count)classFibIterator:def__init__(self,max_count):self.max_countmax_count self.a,self.b0,1self.count0def__next__(self):ifself.countself.max_count:raiseStopIteration self.count1self.a,self.bself.b,self.aself.breturnself.a这样每次for循环都会创建新的FibIterator可以反复遍历。但说实话日常开发中很少需要写两个类用生成器函数更省事deffibs(max_count):a,b0,1for_inrange(max_count):a,bb,abyielda生成器函数自动实现了迭代器协议yield就是__next__的返回值函数结束自动抛StopIteration。除非你需要维护复杂的状态或者实现双向迭代否则别手写迭代器类。itertools 混用技巧itertools是迭代器操作的瑞士军刀但很多人只会用chain和cycle。实际工作中islice和takewhile才是高频利器。比如从一个大文件里读前100行用islice切片迭代器避免一次性加载全部fromitertoolsimportislicedefread_large_file(path):withopen(path)asf:# 这里踩过坑islice 返回的是迭代器不会立即执行forlineinislice(f,100):process(line)再比如处理流式数据时用takewhile按条件截断fromitertoolsimporttakewhiledefprocess_stream(stream):# 遇到负数就停止别用 filter 因为 filter 会遍历全部foritemintakewhile(lambdax:x0,stream):do_something(item)itertools.tee也是个好东西能把一个迭代器克隆成多个独立副本。但注意tee会缓存中间数据如果两个副本遍历速度差距大内存会暴涨。我见过有人用tee处理百万级数据结果内存爆了——后来改成用列表缓存才解决。迭代器与生成器的性能陷阱生成器虽然省内存但每次yield都有上下文切换开销。如果你在循环里频繁调生成器函数性能可能不如列表推导式。实测过生成100万个整数生成器比列表慢约30%但内存占用只有1/10。取舍看场景。另一个坑生成器只能遍历一次。如果你需要多次遍历要么转成列表要么重新创建生成器。我习惯在函数文档里注明“返回生成器只能遍历一次”避免同事误用。实战写一个可重置的迭代器有时候需要迭代器能重置到初始状态比如分页请求失败后重试。可以用itertools.cycle配合islice实现fromitertoolsimportcycle,isliceclassResettableIterator:def__init__(self,data):self.datadata self._cyclecycle(data)self._pos0def__iter__(self):returnselfdef__next__(self):ifself._poslen(self.data):raiseStopIteration self._pos1returnnext(self._cycle)defreset(self):self._cyclecycle(self.data)self._pos0但说实话这种设计有点别扭。更好的做法是用itertools.tee或者直接存列表。我后来重构时干脆把数据存成列表需要迭代时用iter()生成新迭代器——简单可靠。个人经验建议能用生成器就别手写迭代器类。生成器函数自动处理了StopIteration代码量少一半可读性高。只有需要实现__next__之外的接口比如send、throw时才考虑类。调试迭代器时用list()转成列表看内容。别直接print(iterator)那只会打印对象地址。我习惯在__next__里加个临时print确认每次返回的值对不对。itertools的chain和zip配合迭代器时注意长度不一致的问题。zip默认按最短的截断zip_longest会填充None。选哪个取决于业务逻辑别想当然。永远不要在__next__里捕获StopIteration。如果你在迭代器内部捕获了它外部循环就永远收不到终止信号。我见过有人为了“优雅处理”在__next__里try-except吞掉异常结果循环跑飞了。性能敏感场景用for循环而不是whilenext()。for循环底层是 C 实现的迭代协议比 Python 层手动调next()快一个数量级。除非你需要手动控制迭代节奏比如跳过某些元素否则别自己写循环。迭代器协议看着简单但实际用起来坑不少。记住核心__iter__返回迭代器__next__返回元素或抛异常for循环帮你处理了所有脏活。下次遇到迭代器相关的问题先检查这三个点八成能定位到问题。