pythonでイテレータをイテレータのチャンクに分割するコードです。リストをチャンクに分割するコードはネットでいくつか見つけましたが、イテレータをチャンクに分割するコードが見つからないので書いてみました。
最初からリストであるものをチャンクに分割するなら、インデックスを使えば良いし、イテレータをリスト化しつつ分割するのも簡単です。ただ、イテレータをイテレータのまま分割するためには、サブイテレータからベースとなるイテレータに対してStopIterationを通知しなければならず、この方法に苦労しました。結局、ちょっと気持ち悪い方法を使っています。
何が気持ち悪いかというと、subiterの中から外に値を渡すために、exhausted = [False]というリストを使っていることです。普通にbool値として持たせると、中から外に書き込めないため、可変オブジェクトであるリストを使っています。pythonの公式ドキュメントの、itertools.groupbyの参考実装だと似たような状況でclassを使っているので、その方がpythonらしい実装なのかもしれません。
リスト化すると巨大になりすぎ、パフォーマンスに影響するようなケースでも気軽に使えるのが特徴です。
def chunk(it, n):
it = iter(it)
exhausted = [False]
def subiter(it, n, exhausted):
for i in range(n):
try:
yield next(it)
except StopIteration:
exhausted[0] = True
while not exhausted[0]:
yield subiter(it, n, exhausted)