ジェネレーター関数 yield — 値を 1 つずつ生成してメモリを抑える

前回はクロージャで、関数の中に状態を閉じ込めて毎回違う値を返す書き方を見ました。Python にはこのような「呼ぶたびに次の値を返す」関数を専用で作るしくみがあり、それがジェネレーター関数です。returnの代わりにyieldを使うと、関数の途中で処理を止めて値を返し、次に呼ばれたときに続きから再開できます。

大きなデータを一気にメモリに展開せず1 件ずつ流す用途に向いており、ログ処理・大量データの処理などで活躍します。

yield の基本 — 1 つずつ値を返す関数

関数の中にyield 値を書くと、その関数はジェネレーター関数です。普通の関数のように呼び出しても中身は実行されず、generatorという特殊なオブジェクトが返ってくるだけです。

値を実際に取り出すにはnext(オブジェクト)を呼びます。

最初のnext()で関数の冒頭から最初のyieldまで実行され、その値が返されます。

次のnext()でその続きから次のyieldまで進みます。

yieldがもう無くなった状態でさらにnext()を呼ぶと、StopIterationという例外が発生します。

def simple():
    yield 1
    yield 2

gen = simple()
print(type(gen))    # <class 'generator'>

print(next(gen))    # 1
print(next(gen))    # 2
print(next(gen))    # 3
# print(next(gen))  ← StopIteration（もう yield が無い）

return との違い

通常の関数はreturnした時点で関数のスコープごと消え、次に呼んでも最初から実行されます。一方、ジェネレーター関数はyieldで処理を中断し、ローカル変数や進行位置をそのまま保持します。次のnext()で続きから再開できるのが大きな違いです。

注文 ID を 1, 2, 3 と順番に発行するジェネレーター関数を作って、next()で 1 つずつ取り出します。

① def order_ids():を定義し、yield 1 / yield 2 / yield 3を 3 行並べてください。

② gen = order_ids()でジェネレーターオブジェクトを作成してください。

③ print(next(gen))を 3 回続けて呼び出し、1 / 2 / 3と返ってくることを確認してください。

（正しく実行できれば解説が表示されます）

for で順に取り出す — StopIteration を意識しなくていい

next()を毎回書くのは面倒で、StopIterationを自分で処理する必要もありません。代わりにfor value in ジェネレーター:と書けば、Python が裏でnext()を呼び続け、最後まで来たら自動でループを抜けてくれます。普段はこの書き方が圧倒的に多いです。

def count_up_to(max_value):
    print("ジェネレーター開始")
    for i in range(max_value):
        print(f"  yield 直前: {i}")
        yield i
        print(f"  yield 後の続き: {i}")

for v in count_up_to(3):
    print(f"受け取った値: {v}")

# 出力の流れ:
# ジェネレーター開始
#   yield 直前: 0
# 受け取った値: 0
#   yield 後の続き: 0
#   yield 直前: 1
# ...（続く）

顧客データを 1 件ずつ流すジェネレーターをforで受け取って表示します。

① def each_customer():を定義し、for name in ["佐藤", "鈴木", "高橋"]:でループしながらyield nameしてください。

② for name in each_customer():で取り出し、print(f"次のお客様: {name}")を表示してください。

3 件すべてが順に表示されれば成功です。

list との違い — メモリ使用量を抑える

0 から 999_999 までを扱うとき、リスト内包表記でリストを作ると100 万個の整数を一度にメモリに載せることになります。一方、ジェネレーターなら現在の 1 つだけをメモリに置き、必要になった時点で次を計算します。リストはおよそ数 MB のサイズになりますが、ジェネレーターオブジェクト自体は数百バイトしかありません。

import sys

MAX = 10 ** 6

# リスト: 全件を一気にメモリに載せる
data_list = [i for i in range(MAX)]
print(sys.getsizeof(data_list))
# 例: 8000056（およそ 8 MB）

# ジェネレーター式: 現在の 1 件だけ保持
data_gen = (i for i in range(MAX))
print(sys.getsizeof(data_gen))
# 例: 200 程度のバイト数

# 流して使う側のコードは同じ
for v in data_gen:
    if v > 2:
        break
    print(v)
# 0
# 1
# 2

ジェネレーター式という近道

リスト内包表記の角括弧[ ... ]を丸括弧( ... )に変えるだけで、ジェネレーター式になります。(i for i in range(1_000_000))のように書くと、ジェネレーター関数をdefで書かなくても 1 行で同じ効果が得られます。sum() / max() / any()などにそのまま渡せて便利です。

ジェネレーター式（リスト内包表記の角括弧を丸括弧に変えた書き方）で価格データを作り、type()で型を確認してからforで 1 件ずつ取り出します。

① prices = (base * 100 for base in range(1, 6))を作ってください（[ ]ではなく( )を使うのがポイント）。

② print(type(prices))で型を表示し、<class 'generator'>と出ることを確認してください。

③ for p in prices:で取り出して、print(f"価格: {p}")を表示してください。

yield from でジェネレーターを連結する

ジェネレーターの中から別のジェネレーターの値をそのまま流したいときは、for v in sub: yield vのように書く代わりにyield from サブジェネレーターと 1 行で書けます。複数のデータ源を 1 つのジェネレーターにまとめたいときに便利です。

例えば、東京支店と大阪支店、それぞれの売上を流す関数があるとき、yield from tokyo_sales()とyield from osaka_sales()を続けて書けば、呼び出し側からは 1 本の流れに見えるジェネレーターになります。

def tokyo_sales():
    yield 1200
    yield 980

def osaka_sales():
    yield 850
    yield 1340

def all_sales():
    yield from tokyo_sales()
    yield from osaka_sales()

for amount in all_sales():
    print(amount)
# 1200
# 980
# 850
# 1340

支店別の在庫リストを 1 つにまとめて流すジェネレーターを作ります。

① def store_a():でfor item in ["りんご", "みかん"]: yield itemしてください。

② def store_b():でfor item in ["バナナ", "ぶどう", "いちご"]: yield itemしてください。

③ def all_items():を定義し、yield from store_a()のあとにyield from store_b()を続けて書いてください。

④ for item in all_items(): print(item)で 5 件を順に表示してください。