写 Python 写了几年我越来越觉得“迭代器”是那种看似简单、用起来才知道深浅的概念。有一个需求我做了很多遍从很大的数据源里按批次取数喂给下游逻辑处理。最早我写的是“先把整个文件读进列表再按切片一组一组取”文件小的时候一切正常文件一旦上了几个 GB内存直接告急。后来把逻辑都改成了自定义迭代器这个问题才算真正解决。迭代器本质上是一个“按需生产数据”的机制但把它设计得顺手、可控、不踩坑需要好几层思考。这篇就围绕自定义迭代器设计展开从协议原理到大数据批量加载的实战写法再到练习关卡里的常见陷阱把我在项目中反复用到的设计套路一次说清楚。1. 内置迭代器不够用的时候先想清楚你要解决的四个问题1.1 内存压力整块加载 vs 按需取用Python 的内置迭代能力很强range、list、字典、文件对象都能直接for x in data。但对“数据批量加载”这个场景内置方案往往不是最优解。最典型的例子就是读大文件# 很多人习惯的第一版写法 with open(train.txt, r, encodingutf-8) as f: lines f.readlines() for line in lines: process(line)readlines()会把所有行一次性塞进内存。文件 2GB 的时候这个列表加上 Python 字符串对象的开销实际内存占用可能跑到 5GB 甚至更高。就算机器内存够大这个写法在数据量持续增长时也不可持续。改成文件对象直接迭代with open(train.txt, r, encodingutf-8) as f: for line in f: process(line)内存占用立刻降下来了因为文件对象本身就是内置的迭代器每次next()只读取一行。但问题也随之而来我想“一批一批”地处理比如每 100 行合成一个 batch 再喂给模型每次攒够 100 行就把列表交出去。这个过程如果自己在 for 循环里用临时列表收着代码会非常啰嗦而且批大小的控制逻辑会散落在各处。这时候自定义迭代器的第一个价值就出来了把“按需取数”的规则封装成一个对象调用方只负责消费不关心内部怎么攒批、怎么判断结束。1.2 状态管理让数据游标跟着业务走内置迭代器的问题不只是内存。for line in f这种写法把迭代状态藏在文件对象内部一旦循环中间break了这个文件对象处于“读到一半”的状态。如果你需要多种消费模式比如先读头部 100 行做统计再从头开始按批读取内置迭代器会让你很难办——因为迭代器是一次性的走完就没了除非重新打开文件。自定义迭代器可以把“游标位置”“是否已读完”“当前批大小”这些状态收拢到一个类里。比如我可以设计一个对象既能按批读取也能通过reset()回到文件头。对于下游逻辑来说它面对的是一个稳定的接口每次next()给我一批数据当我说“重新来一次”时它知道该怎么做。这就是状态封装的价值。1.3 一次性消费的控制权内置迭代器还有一个隐藏属性只能向后不能回头。list(iterator)之后原来的迭代器就空了。但在真实业务里你可能需要对同一个数据源做多次扫描比如训练时每个 epoch 都要重新读取一遍全部数据。如果直接用内置迭代器你得每次重新创建一个新的迭代器对象。自定义迭代器可以设计成“可重入”让同一个对象支持反复迭代这对模型训练这类场景特别重要。1.4 多消费者共享同一数据流还有一种场景是多个消费者需要共享同一份数据流。比如一份日志文件我希望一个消费者统计错误率另一个消费者解析特定字段。如果各自打开文件不仅浪费 IO还可能因为缓存不一致看到不同状态。自定义迭代器可以作为“数据分发中心”内部维护一个游标多个消费者通过同一实例拿数据可以保证大家看到的是同一顺序、同一份数据。这个需求用内置迭代器做会比较别扭但用自定义迭代器加一点内部逻辑就能实现。2. 先把协议吃透for循环到底对你的对象做了什么2.1 最小迭代器类五行代码看穿协议在动手设计复杂迭代器之前先把最核心的协议搞明白。Python 的迭代协议实际上只有两个方法__iter__()返回一个迭代器对象。__next__()每次被调用时返回下一个元素没有元素时抛出StopIteration。一个最基本自定义迭代器长这样class Counter: def __init__(self, limit): self._n 0 self._limit limit def __iter__(self): return self def __next__(self): if self._n self._limit: raise StopIteration self._n 1 return self._n然后就可以这样使用for num in Counter(3): print(num) # 输出1 2 3这个类很小但它包含了自定义迭代器的全部关键点__iter__返回自身__next__负责状态推进和结束判定。for循环内部做了三件事先调用iter(obj)拿到迭代器然后不断调用next(iterator)直到捕获到StopIteration就结束循环。2.2 可迭代对象与迭代器的区别很多初学者会把“可迭代对象iterable”和“迭代器iterator”混为一谈。列表、字典、集合都是可迭代对象但它们本身不是迭代器。你调用iter([1, 2, 3])会返回一个列表迭代器这个迭代器才实现真正的__next__。而[1, 2, 3]自己并不实现__next__。验证一下就清楚了from collections.abc import Iterator, Iterable print(isinstance([], Iterable)) # True print(isinstance([], Iterator)) # False print(isinstance(iter([]), Iterator)) # True这个区别直接影响了自定义迭代器的设计。如果你的类同时实现__iter__和__next__那它既是可迭代对象也是迭代器。这时候__iter__返回self没问题因为self已经满足迭代器接口。如果你的类只实现__iter__返回一个新的迭代器对象那它就是“可迭代对象但不自己当迭代器”比如列表就是这种模式。想清楚你设计的是哪一种后面才不会在“重复迭代”这个问题上栽跟头。2.3 __iter__返回self的适用边界当__iter__返回self时意味着同一个对象只能被完整遍历一次。因为迭代状态比如self._n是对象内部的一个字段第一次for循环把它推进到终点后第二次for循环拿到的还是同一个对象状态已经耗尽循环不会执行。如果你希望同一个对象能被多次遍历就不要在__iter__里返回self而是返回一个新创建的状态副本。比如class ReusableCounter: def __init__(self, limit): self._limit limit def __iter__(self): return CounterIterator(self._limit)CounterIterator是另一个类保存当前计数状态。这样每次for循环都会拿到一个新的迭代器互不干扰同一个ReusableCounter对象可以被反复遍历。像这种“迭代器对象”和“可迭代容器对象”分离的设计在数据加载场景里非常有用——数据源是同一个但每次训练 epoch 都从头开始。第 3 节里的批量加载器会同时展示这两种模式。3. 实战为大数据加载场景设计一个可复用的批量迭代器3.1 需求拆解与设计决策放到真实项目里自定义迭代器最典型的用途就是“数据批量加载”。假设我有一个 CSV 文件每行是一条样本下游模型一次要吃 100 条。需求拆开看有四个关键点不能一次性把整个文件读入内存。每次next()返回一个包含 100 行数据的列表。文件读取结束后自动抛出StopIteration。对象本身可以“重置”便于多次遍历。第一个版本我选择让__iter__返回self因为顺序消费是主要场景实现简单状态管理直观。同时提供一个reset()方法负责把文件游标移回开头。3.2 代码实现BatchFileReaderclass BatchFileReader: def __init__(self, file_path, batch_size100): self._file_path file_path self._batch_size batch_size self._fp None self._opened False def _open_if_needed(self): if self._fp is None: self._fp open(self._file_path, r, encodingutf-8) self._opened True def __iter__(self): self._open_if_needed() return self def __next__(self): self._open_if_needed() batch [] for _ in range(self._batch_size): line self._fp.readline() if not line: break batch.append(line.rstrip(\n)) if not batch: raise StopIteration return batch def reset(self): if self._fp is not None: self._fp.seek(0)使用方式非常直接reader BatchFileReader(train.csv, batch_size100) for batch in reader: process(batch) # 想再来一轮 reader.reset() for batch in reader: process(batch)这个实现有几个细节值得解释。_fp.readline()返回空字符串时表示文件读完了这就是“读完”信号。注意if not line这种判断对最后一行没有换行符的文件依然有效不会漏数据。攒了一批数据但发现是空列表说明文件已经到了末尾此时必须抛StopIteration否则for循环会无限运行。3.3 边界情况与异常处理实际使用中文件迭代器最容易被忽略的是资源释放。我见过很多同学写完自定义迭代器类里打开了文件却没有地方关闭最后靠垃圾回收救场。上面这个版本有一个隐患如果外部for循环提前breakStopIteration可能还没等到文件句柄就一直开着。改进方案是让BatchFileReader实现上下文管理器或者把文件打开逻辑放到__enter__/__exit__里class BatchFileReader: def __init__(self, file_path, batch_size100): self._file_path file_path self._batch_size batch_size self._fp None def __enter__(self): self._fp open(self._file_path, r, encodingutf-8) return self def __exit__(self, exc_type, exc_val, exc_tb): if self._fp: self._fp.close() return False def __next__(self): if self._fp is None: raise RuntimeError(请在 with 语句中使用 BatchFileReader) batch [] for _ in range(self._batch_size): line self._fp.readline() if not line: break batch.append(line.rstrip(\n)) if not batch: raise StopIteration return batch这样把“打开文件”的责任明确交给调用方配合with语句资源安全就兜住了。还有一个小边界文件为空第一行readline()就返回空字符串那么第一次调用next()时batch为空直接抛StopIterationfor循环一次都不执行逻辑是正确的。with BatchFileReader(train.csv, batch_size100) as reader: for batch in reader: process(batch)3.4 把骨架抽出来可复用迭代器基类连续写过几个批量加载迭代器之后我发现它们的主干结构是完全一致的推进一次、判断是否结束、返回一个数据单元。区别只在于“下一次数据从哪里来”。于是我把骨架抽成一个基类后续新数据源只需要实现一个方法class BaseStreamIterator: def __iter__(self): return self def __next__(self): item self._fetch_next() if item is None: raise StopIteration return item def _fetch_next(self): raise NotImplementedError注意这里约定_fetch_next返回None表示数据源已经枯竭。实际业务里数据单元本身不太可能是None所以这个约定是安全的。如果是数据库游标就写一个DatabaseCursorIterator如果是队列就写一个QueueIterator。每个子类只需要关心“拿到下一条”迭代边界和循环控制逻辑都复用基类的。这比每个项目从零写一套__next__要省心得多排查问题也集中。4. 同一个批加载需求用生成器写会更短但未必更合适4.1 生成器版本十行代码搞定同款逻辑如果只是想尽快搞定批量加载Python 的生成器语法比自定义迭代器类短得多。同一个BatchFileReader功能用生成器写出来是def batch_file_reader(file_path, batch_size100): with open(file_path, r, encodingutf-8) as fp: while True: batch [] for _ in range(batch_size): line fp.readline() if not line: break batch.append(line.rstrip(\n)) if not batch: break yield batch调用时甚至不需要创建对象直接for batch in batch_file_reader(train.csv, 100)。带yield的函数调用后会返回一个生成器对象它天然实现了迭代器协议。文件句柄在with块中自动管理break之后上下文管理器也会把文件关掉。从代码量、资源安全、可读性三个维度看生成器版本都更占优。4.2 类与生成器的取舍状态可检查、可重入、可双向控制那为什么我还要花力气写自定义迭代器类因为生成器在一个关键维度上比较弱它不具备可检查的外部状态接口。生成器的内部状态完全封闭在函数栈帧里外部只能一个接一个地next()想看当前游标在哪、还剩多少数据、是否已经读完都没有入口。以批量加载为例训练脚本经常需要知道“当前轮次已经产出了多少个 batch”或者“这个 epoch 还剩多少样本”。用类实现时我可以在类上维护self._batch_count随时用reader.batch_count查看。生成器要实现同样的能力要么额外包一层统计器要么用闭包维护可变状态代码会变复杂。类方案还有一个优势是“双向控制”。生成器虽然支持send()和throw()但写复杂了可读性很差。而类迭代器可以把暂停、恢复、跳过的逻辑写成普通方法。比如某个迭代器需要支持“跳过前 N 条”或者“当前批次异常时重试”类方案可以直接加方法class BatchFileReader: def skip(self, n): for _ in range(n): if not self._fp.readline(): raise StopIteration生成器想实现同样的操作只能在yield之后接收外面发来的指令代码写起来像是“为了双向通信而双向通信”不如类直接。类型上的区别也值得说清楚。生成器对象既是可迭代对象也是迭代器但它同样是一次性的。这批数据消费完生成器自动进入StopIteration再想从头来必须重新调用函数生成一个新对象。如果业务上确实需要同一个数据源反复遍历类方案可以在__iter__里返回新的内部迭代器从而做到“源对象可复用”。这是很多人在“生成器 vs 自定义迭代器”里容易忽略的分界线。4.3 生成器的隐藏限制无法倒带状态封闭再补一个实战中踩过的坑。生成器函数里with open(...) as fp打开文件文件句柄的生命周期绑在生成器的生命周期上。如果一个消费者只取了前 3 个 batch 就把生成器丢弃了文件在生成器被垃圾回收前不会及时关闭。虽然 CPython 的引用计数机制通常能快速回收但依赖垃圾回收来管理文件句柄总归是悬着的。类方案借助__del__或with上下文管理器可以更清晰地控制释放时机。还有一点生成器不能中途被“检查内部状态”一旦你在调试时需要知道“当前读到文件哪一行”生成器版本只能靠外部计数类版本直接读属性就行。这种差异在业务代码里可能只是调试便利在框架设计层面就可能是选型依据。我的习惯是一次性消费、代码量优先用生成器需要复跑、需要观察进度、需要支持复杂控制逻辑时用自定义迭代器类。5. 练习关卡的经典坑StopIteration胡乱抛出与迭代器“一次消费”的陷阱5.1 不抛StopIteration会怎样for循环直接变成死循环很多 Python 学习课程里都会有类似“第 4 关迭代器”的关卡要求实现一个自定义迭代器类。我在帮人 review 代码时发现最常见的错误就是__next__在数据耗尽时不抛StopIteration而是返回None或直接return。看这个class BrokenCounter: def __init__(self, limit): self._n 0 self._limit limit def __iter__(self): return self def __next__(self): if self._n self._limit: return None # 错误示范 self._n 1 return self._n你用list(BrokenCounter(3))试试程序会永远转下去。因为for循环结束的唯一信号就是StopIteration异常如果next()返回None循环会认为“又拿到一个新元素”继续推进。别小看这个坑代码一旦写成这样表现是卡死而不是报错排查起来特别费劲。正确的做法是def __next__(self): if self._n self._limit: raise StopIteration self._n 1 return self._nPython 3.7 之后生成器有个便利保证生成器函数正常执行完会自动抛出StopIteration所以写生成器时不需要手动抛。但自定义迭代器类没有这个保证你必须手动写清楚。5.2 迭代器与可迭代对象的混淆list一次和两次的结果练习关卡里还有一个高频考点为什么同一个迭代器第二次list()是空列表。看这段代码it Counter(5) print(list(it)) # [1, 2, 3, 4, 5] print(list(it)) # []因为Counter的__iter__返回selflist(it)第一次就把内部状态推到了_n 5第二次list(it)再调用next()时直接碰到StopIteration。这个行为对“迭代器”来说是符合协议的但对不熟悉的人来说非常反直觉。这才是第 2 节提到的“可迭代对象 vs 迭代器”概念带来的实际问题。如果你希望list(x)两次结果一致那么x应该是一个可迭代对象而不是一个迭代器。比如list([1,2,3])两次结果一样因为列表本身不是迭代器每次iter(list)会生成新的列表迭代器。所以在设计练习答案或业务代码时先问自己一句这个对象是想表现“容器”的语义还是“流”的语义容器语义每次迭代独立流语义一次消费完毕。最怕的是两者混着设计外部使用者没法预期行为。5.3 通关题里最常见的三类需求模板从“第 4 关迭代器”这类练习里提炼一下我发现大多数关卡其实是在考三个模板第一类是“有限序列迭代器”比如生成 1 到 N 的整数。核心是在__next__里判断边界并抛StopIteration。第二类是“无限序列迭代器”比如斐波那契数列。此时没有边界判断next()永远能返回新数外部要靠itertools.islice截断class FibonacciIterator: def __iter__(self): self._a, self._b 0, 1 return self def __next__(self): self._a, self._b self._b, self._a self._b return self._a from itertools import islice for num in islice(FibonacciIterator(), 8): print(num)第三类是“批量聚合迭代器”也就是本文主打的场景。它和上面两类最大的不同在于真正的数据并不是一次一条而是攒一批才吐一次。边界条件从“单个元素为空”变成了“整个批次为空”。把这三个模板吃透绝大多数迭代器考题就都没问题了而且能直接迁移到生产代码里。6. 结束前分享三个我反复使用迭代器的真实项目片段6.1 批量训练时的数据喂养器做模型训练时我的数据加载管道是这样组织的底层是用BatchFileReader按批读文件中间层用itertools.cycle把有限数据流变成无限循环供训练循环每次取一个 batch。因为BatchFileReader实现了reset()配合itertools.cycle时整个流程非常顺畅。关键收益是内存占用始终是“一批数据”的量级而不是“整个数据集”的量级。500GB 的训练语料也能在普通单机上跑完整流程这在用readlines()的年代是不可想象的。6.2 带指数退避的无限重试迭代器另一个项目需要持续从一个外部接口拉取消息失败时按指数退避重试。拉取服务本身就是一个自定义迭代器每次next()拿一条消息内部如果发现连接断开自动按 1 秒、2 秒、4 秒的节奏重连重连成功后继续产出。调用方完全无感它只知道“只要我不主动停止迭代数据就会一直来”。用类实现时重试状态可以清晰放在属性里调试时能直接看到当前退避到第几轮。这个逻辑如果硬塞进生成器里yield中间夹着重连代码可读性会差很多。6.3 自定义迭代器与itertools的组合玩法最后一个小技巧把自定义迭代器当成“基础零件”上层用itertools来拼出各种复杂逻辑。比如takewhile控制读取条数islice跳过中间段tee把同一数据流复制给多个消费者并行处理。这些都是标准库自带的配合自定义迭代器使用效果很好。我常用的组合是from itertools import islice reader BatchFileReader(train.csv, batch_size100) first_1000 islice(reader, 10) # 前 10 个 batch for batch in first_1000: process(batch)islice消费完前 10 个 batch 后reader 内部的文件游标恰好停在“第 1000 条样本之后”数据源本身没有被打乱。这种“外部裁剪 内部维持游标”的组合优势正是自定义迭代器状态封装带来的。如果换成裸文件对象或裸生成器想精确跳过前 1000 条样本再继续读你得自己在循环里数数又绕又容易错。写到这里自定义迭代器设计这件事基本讲透了。核心还是那两点理解迭代器协议想清楚自己要“容器语义”还是“流语义”。具体到代码层面能复用生成器就复用生成器需要状态管理、反复遍历、精细控制时再上类。希望这篇把我在批量加载、无限重试、评估练习里攒下来的经验说明白了下次你再被数据撑爆内存或者被“迭代器为什么第二次是空的”难住时能有个清晰的方向。