【python】generator を「チャンク化された generator の generator」に変換する小技

1. はじめに

Python で巨大なデータやストリームを扱うとき、
「イテレータを一定サイズごとに区切って処理したい」
という場面がある。

例えば:

• 大きなファイルを N 行ずつ処理したい
• generator をメモリに載せず、バッチ単位で流したい
• for 文を二重ループにしたいが、list 化は避けたい

といったケースだ。

本記事では、
generator を「size 個ずつ区切られた generator の generator」に変換する
シンプルかつ安全な実装を紹介する。

2. 実現したい挙動

以下のようなコードを書けるようにしたい。

for (i,chunk) in enumerate(get_chunk(range(10),3)):
    print(f"{i=}, {chunk=}")
    for j in chunk:
        print(f"{j} ", end="")
    print("\n---")

出力

i=0, chunk=<generator object _subgen at 0x00000201718B1490>
0 1 2 
---
i=1, chunk=<generator object _subgen at 0x00000201718B1BE0>
3 4 5 
---
i=2, chunk=<generator object _subgen at 0x00000201718B1490>
6 7 8 
---
i=3, chunk=<generator object _subgen at 0x00000201718B1BE0>
9 
---

ポイントは:

• chunk 自体が generator

• 最後のチャンクは size 未満でも返る

• 元の generator を一切 list 化などしていない
  → メモリに一度に載せる必要がない

という点だ。

3. 実装

コードは以下の通り。

get_chunk.py

from itertools import islice

def get_chunk(gen, size):
    """
    generator を size 毎に区切った
    generatorのgenerator に変換する
    """
    it = iter(gen)
    while True:
        chunk = islice(it, size)
        try: first = next(chunk)
        except StopIteration: return
        yield _subgen(first, chunk)

def _subgen(first, rest):
    yield first
    yield from rest

3-1. islice だけでは足りない理由

よくあるのが:

while True:
    chunk = list(islice(it, size))
    if not chunk: break
    yield chunk

という実装である。

listや tuple など の generator を作りたい場合はこれでよいが、
今回は generator の generator を作りたいので適さない。

蛇足だが、 tuple の generator でよいのであれば、
Python 3.12 以降なら itertools.batched を利用して:

from itertools import batched

for chunk in batched(range(10), 3):
    print(chunk)

で十分だ。

3-2. first = next(chunk) の意味

islice は空でも generator としては存在する。
その場合、空の generator が yield されてしまう。

だが今回は、空の場合は yield させたくないので、
next(chunk) が成功することにて空でないことを確認している。

空の場合、 next(chunk) が失敗して StopIteration 例外が発生するので、
その場合には直ちに return することで、 yield を避けている。

3-3. なぜ _subgen を使うのか

first = next(chunk) で 1 要素を消費してしまうため:

yield chunk

とすると先頭要素 first が欠落してしまう。

そこで:

yield _subgen(first, chunk)

と書き、 first と 残りの chunk を統合して yield している。
_subgen では:

サブジェネレータ _subgen

yield first
yield from chunk

という処理をすればよい。

3-4. この実装のメリットまとめ

• 完全な遅延評価（lazy）
  • メモリ消費がほとんどない（list 不使用）
• 最後の余り要素も自然に処理可能
• for × for の構造が直感的に書ける

特に 巨大データ処理・ストリーム処理・memmap / database cursor との相性が非常に良いだろう。

4. まとめ

まとめ

「generator を chunk に分けたい」という問題は一見単純だが:

• 空チャンク
• 要素の取りこぼし
• 不要な list 化

といった落とし穴が多くある。

今回の実装はそれらをすべて回避しつつ、
短く・安全で・読みやすい 形にまとめた。

必要な場面では、ぜひそのままコピペして使ってほしい。