pythonでイテレータをチャンクに分割する

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)