モダンPython中級 #4 イテラブル / ジェネレータ / yield from

プログラミング言語 Python generator iterable
読了 7分

基礎 #4 の最後で軽く触れたジェネレータ式 (x for x in iter) — 今回がそのテーマです。for がどう動くのか から始めて、ユーザー定義のイテラブル、ジェネレータ関数、yield from まで扱います。

for in の正体 — イテラブルプロトコル #

よく使う形 
for x in [1, 2, 3]:
    print(x)

この一行が内部で何をしているのか紐解くと:

実際の流れ 
items = [1, 2, 3]
it = iter(items)        # 1) イテラブル → イテレータ
while True:
    try:
        x = next(it)    # 2) 次の値を要求
    except StopIteration:
        break           # 3) 終わったら止まる
    print(x)

要点となる 2 ステップ:

  1. iter(obj) — イテラブルからイテレータを得る (__iter__ を呼ぶ)
  2. next(it) — 次の値を要求 (__next__ を呼ぶ)、終わっていれば StopIteration を投げる

イテラブル vs イテレータ #

用語が混乱するところなので整理します。

定義メソッド例
イテラブル (Iterable)iter() ができるすべてのもの__iter__list、dict、str、range、ファイル、ジェネレータ
イテレータ (Iterator)「次の値」を返せるもの__next__ (と __iter__)iter([1,2,3]) の結果、ジェネレータ

すべてのイテレータはイテラブル です (自分自身を返す __iter__ を持つ)。逆は成り立ちません — list はイテラブルですがイテレータではありません。next(my_list) はエラーになります。

ユーザー定義のイテラブル — クラスで #

Range を自作 
class MyRange:
    def __init__(self, start: int, stop: int):
        self.start = start
        self.stop = stop

    def __iter__(self):
        return MyRangeIterator(self.start, self.stop)

class MyRangeIterator:
    def __init__(self, current: int, stop: int):
        self.current = current
        self.stop = stop

    def __iter__(self):
        return self

    def __next__(self):
        if self.current >= self.stop:
            raise StopIteration
        value = self.current
        self.current += 1
        return value

for x in MyRange(0, 3):
    print(x)
# 0, 1, 2

2 つのクラス — イテラブルとイテレータを分けます。イテラブルは 何度も巡回できます が、イテレータは一度使い切ったら終わりです。

何度も巡回できる理由 
r = MyRange(0, 3)
list(r)    # [0, 1, 2]
list(r)    # [0, 1, 2]  ← 毎回新しいイテレータ

ジェネレータ関数 — 同じことを関数ひとつで #

上のコード 2 クラスを yield を含む関数ひとつ に縮められます。

ジェネレータ関数 
def my_range(start: int, stop: int):
    current = start
    while current < stop:
        yield current
        current += 1

for x in my_range(0, 3):
    print(x)
# 0, 1, 2

関数本体に yield が一度でも入っていれば、その関数は呼び出すと 通常の値を返すのではなく、ジェネレータオブジェクトを返します。上のクラス 2 つと同じことをします。

yield の動き方 #

最も混乱する部分がこれです。

ジェネレータを作る 
def gen():
    print("step 1")
    yield 1
    print("step 2")
    yield 2
    print("step 3")

g = gen()
# ここまでで関数本体はまだ実行されていない!

g = gen() だけでは 関数本体は実行されません。 最初の next(g) が来てから始まります。

値を要求 
print(next(g))
# step 1
# 1

print(next(g))
# step 2
# 2

print(next(g))
# step 3
# StopIteration ← yield がもうないと

各 yield で関数が 一時停止 します。次の next() 呼び出しが来ると、そこから再開します。関数の実行時間が呼び出し側とインターリーブされる のがジェネレータの核心です。

ジェネレータ式とどう違うか #

基礎 #4 で見た (x for x in iter) と同じ動作をしますが、関数の形の方が表現力が大きいです。

ジェネレータ式 vs 関数 
# 一行で表現可能 — 式が合う
squares = (x ** 2 for x in range(10))

# 複雑なロジック — 関数が合う
def squares_evens_only():
    for x in range(10):
        if x % 2 != 0:
            continue
        yield x ** 2

遅延の価値 — メモリと速度 #

ジェネレータの最大の価値は すべての値を一度に作らない という点です。

100 万個を処理 
# リスト内包表記 — 100 万個を即時生成、メモリを全部使う
squares_list = [x ** 2 for x in range(1_000_000)]

# ジェネレータ — 必要なときだけ生成、メモリはほとんど使わない
squares_gen = (x ** 2 for x in range(1_000_000))

total = sum(squares_gen)   # 一度の巡回で完了

無限シーケンスも可能 #

無限シーケンス 
def counter(start: int = 0):
    n = start
    while True:
        yield n
        n += 1

# 最初の 5 つだけ使う
from itertools import islice
first_five = list(islice(counter(), 5))
print(first_five)   # [0, 1, 2, 3, 4]

リストでは不可能なことです。ジェネレータは 要求された分だけ 値を作るので無限でも大丈夫です。

パイプライン — ジェネレータをつなぎ合わせる #

各段階がジェネレータのデータ処理パイプラインを作ると、メモリ効率が良く意図が明確になります。

パイプライン 
def read_lines(path: str):
    with open(path) as f:
        for line in f:
            yield line.rstrip()

def filter_errors(lines):
    for line in lines:
        if "ERROR" in line:
            yield line

def parse_timestamp(lines):
    for line in lines:
        ts, _, msg = line.partition(" ")
        yield (ts, msg)

# 合わせて使う
errors = parse_timestamp(filter_errors(read_lines("app.log")))
for ts, msg in errors:
    print(ts, msg)

各段階が 一行ずつだけ処理 します。100GB のファイルでもメモリにすべて載せません。

yield from — ジェネレータの委譲 #

別のイテラブルの値をそのまま流したいとき。

🚫 直接展開して書く 
def chain_two(a, b):
    for x in a:
        yield x
    for y in b:
        yield y
✅ yield from 
def chain_two(a, b):
    yield from a
    yield from b

同じことをしますが yield from の方が短く、それ以外にも 2 つの追加効果があります:

  1. send/throw が自動で委譲される (下で扱う)
  2. 下位ジェネレータの戻り値を受け取れる

ツリー / 再帰の巡回に自然 #

再帰でツリーを平坦化 
def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

result = list(flatten([1, [2, [3, [4]], 5]]))
print(result)   # [1, 2, 3, 4, 5]

yield from flatten(...) の一行が再帰を自然に解いてくれます。

send、throw、close — コルーチンの機能 #

ジェネレータは 値を受け取ることもできます。 yield の結果を変数に受け取れば、外部から send で値を入れられます。

send 
def echo():
    while True:
        received = yield
        print(f"受信: {received}")

g = echo()
next(g)            # 最初の yield まで進める (priming)
g.send("hello")    # 受信: hello
g.send("world")    # 受信: world

このメカニズムの上に非同期 (#7) と協調的マルチタスクが作られています。ただ通常のコードで send を直接扱うことはほぼありません。概念だけ認識 しておけばよいです。

throw と close 
g.throw(ValueError, "例外注入")   # ジェネレータの中で raise と同等
g.close()                          # ジェネレータ終了 (GeneratorExit を投げる)

close() はよく使われます — リソースの後片付けが必要なジェネレータで try/finally に後片付けコードを入れると close 時点で実行されます。

リソースの後片付け 
def read_lines(path):
    f = open(path)
    try:
        for line in f:
            yield line
    finally:
        f.close()

このコードはジェネレータを最後まで回さなくても (break で抜けても) ガベージコレクションのタイミングで close() が呼ばれてファイルが閉じられます。

itertools — 標準ライブラリの宝石 #

データパイプラインでよく使うツールが itertools にあります。

よく使う itertools 
from itertools import (
    count, cycle, repeat,                # 無限
    islice,                               # スライス
    chain,                                # 連結
    groupby,                              # グルーピング
    accumulate,                           # 累積
    combinations, permutations, product,  # 組み合わせ
    starmap, filterfalse, dropwhile, takewhile,  # 変換 / フィルタ
)

# 最初の N 個
list(islice(count(), 5))             # [0, 1, 2, 3, 4]

# 複数のイテラブルを連結
list(chain([1, 2], [3, 4]))          # [1, 2, 3, 4]

# 累積和
list(accumulate([1, 2, 3, 4]))       # [1, 3, 6, 10]

# グルーピング (ソートが必要)
data = [("a", 1), ("a", 2), ("b", 3)]
for key, group in groupby(data, key=lambda x: x[0]):
    print(key, list(group))
# a [('a', 1), ('a', 2)]
# b [('b', 3)]

データパイプラインを書く機会が多いなら一度ざっと見ておくと一生の助けになります。

標準ライブラリのコレクションも同じプロトコル #

collections.abc の ABC たちは このプロトコルを形式化したもの です。

collections.abc — インターフェース 
from collections.abc import Iterable, Iterator

def consume(items: Iterable[int]) -> int:
    total = 0
    for x in items:
        total += x
    return total

# list、tuple、set、ジェネレータ、range、... すべて通る
consume([1, 2, 3])
consume(range(10))
consume(x for x in [1, 2, 3])

関数が受け取る型を Iterable[T] で書けば最も広く安全です。あえて list[T] に絞る必要はありません — 呼び出し側がジェネレータを渡そうがセットを渡そうが同じように動きます。

@contextmanager は実はジェネレータ #

#3 の @contextmanager がどう動くのか、ここで見えてきます。

@contextmanager 再び 
from contextlib import contextmanager

@contextmanager
def chdir(path):
    old = os.getcwd()
    os.chdir(path)
    try:
        yield path
    finally:
        os.chdir(old)

この関数は yield を持つので ジェネレータ関数 です。@contextmanager がそのジェネレータを受け取って __enter__ が最初の next を呼び、__exit__ が二度目の next または throw を呼ぶ オブジェクトに包みます。コンテキストマネージャーがジェネレータの上に作られたツールであることが、ここで確認できます。

まとめ #

今回押さえたこと:

  • for = iter() + next() + StopIteration のシュガー
  • イテラブル (__iter__) ⊃ イテレータ (__iter__ + __next__)
  • ジェネレータ関数 — yield が一度でもあれば関数呼び出しがジェネレータオブジェクトを返す
  • yield ごとに一時停止、次の next で再開
  • 遅延の価値 — メモリ / 無限シーケンス / パイプライン
  • yield from — 別のイテラブルへ委譲、再帰に自然
  • send / throw / close — コルーチンメカニズム、close と try/finally でリソース後片付け
  • itertools 標準ツール集
  • 関数の引数は最も広く Iterable[T] で書こう
  • @contextmanager はジェネレータの上に乗せたツール

次回 (#5 デコレータパターン) では、関数とクラスを包むツール — デコレータ のあらゆるパターンを扱います。@contextmanager、@dataclass もその一形態でした。

X