Pythonのジェネレータとyieldの使い方

最終更新日 2023/10/08

実行時の環境
Python	3.11.5

Pythonのジェネレータは、一言でいえば「イテレータを生成する機能」です。使ったことがなくとも、コードリーディングをしているときにときどき見かけるので、覚えておくと便利な機能となります。
yieldはジェネレータ関数を作成するときに使う「処理を終了せず返り値を返す予約語」のことです。

イテレータは「繰り返して一つずつ取り出せるオブジェクト」です。

実際に手を動かしてその機能を紹介していきます。

yieldとreturnの違い

ジェネレータの説明を始める前に、yieldとreturnの違いを簡単に説明します。

yield 関数の処理を停止して値を返す
return 関数の処理を終了して値を返す

yieldが「停止」となっているのは、「再開」が可能だからです。

ジェネレータ関数を作ってみる

まずは簡単なジェネレータ関数の例を作ってみます。

def sample():
    yield 1
    yield 2
    yield 3

gen = sample()
print(gen)
print(type(gen))
print(next(gen))
print(next(gen))
print(next(gen))

このスクリプトを実行すると、以下のような結果が得られます。

<generator object sample at 0x104bf09e0>
<class 'generator'>
1
2
3

1つ目のジェネレータオブジェクトがいわゆる「繰り返せるオブジェクト（イテラブルオブジェクト）」であり、for文で取り出すことも可能です。
以下でも1, 2, 3の結果が得られます。

def sample():
    yield 1
    yield 2
    yield 3

gen = sample()
[print(i) for i in gen]

List（配列）型へ変換

PythonのList型へ変換したい場合はlist()関数を使います。

def sample():
    yield 1
    yield 2
    yield 3

gen = list(sample())
print(type(gen))
print(gen)

結果は以下です。

<class 'list'>
[1, 2, 3]

ジェネレータにオブジェクトを追加する

一度生成したジェネレータの中にオブジェクトを追加する場合はsend()関数を使います。

def sample():
    v = yield 2
    result = yield v
    yield result

gen = sample()
print(next(gen))
print(gen.send(4))
print(gen.send(8))
print(gen.send(16))

これの出力結果は以下です。

2
4
8
Traceback (most recent call last):
  File "sample.py", line 10, in <module>
    print(gen.send(16))
          ^^^^^^^^^^^^
StopIteration

この処理の流れを説明すると以下のようになります。

ジェネレータ関数を実行する
最初のnext()で2が返される
次の処理からsend()で値が送られた場合、変数vに格納される
さらに次にsend()で値を送った場合、変数resultに格納される
resultまで返し終わった後にsend()しようとしてもエラーが出る

ジェネレータのエラー処理

ジェネレータに内包されるオブジェクトの数に限りがある場合、次の要素を取り出すnext()関数を限界以上に実行するとエラーが出ます。

def sample():
    yield 1
    yield 2
    yield 3

gen = sample()
print(next(gen))
print(next(gen))
print(next(gen))
print(next(gen)) # ここでエラーが出る

結果は以下です。

1
2
3
Traceback (most recent call last):
  File "sample.py", line 10, in <module>
    print(next(gen))
          ^^^^^^^^^
StopIteration

エラーハンドリングしたい場合は以下のように書きます。

def sample():
    yield 1
    yield 2
    yield 3

gen = sample()

for i in range(4):
    try:
        print(next(gen))
    except StopIteration as e:
        print('StopIteration raised')
        break

for文などで繰り返し処理でオブジェクトを処理する場合、このように例外発生でbreakするなどします。

ジェネレータでreturnした場合

ジェネレータ内でreturnを使って値を返した場合、StopIteration例外が発生します。

def sample():
    v = yield 2
    result = yield v

    if result != 4:
        return 'Error' # 文字列でエラー内容を返した場合

    yield result

gen = sample()
print(next(gen))
print(gen.send(4))
print(gen.send(8))

これが出力する結果は以下です。

2
4
Traceback (most recent call last):
  File "/Users/kt2763/tmp/frkz-python-generator/sample.py", line 13, in <module>
    print(gen.send(8))
          ^^^^^^^^^^^
StopIteration: Error

print(gen.send(8))以降の処理はreturnされているため発生しません。

ジェネレータの内包表記

List型と同じく、内包表記でジェネレータを書けます。List型だと[]ですが、ジェネレータの場合は()で囲みます。

squares = (x**2 for x in range(1, 11))

for i in range(5):
    print(next(squares))

これは1 ~ 11の範囲で階乗のジェネレータを作成し、5までの取り出したケースとなります。出力結果は以下です。

ほとんど内包表記で書くことはありませんが、こういった使い方も可能です。

ジェネレータの使い所

今までジェネレータについて紹介してきましたが、「一体いつ使うのか」「List型でよくないか」と感じる人もいることでしょう。
いくつかのユースケースを紹介します。

大きなファイルの中身を取り出すとき

たとえば、巨大なCSVファイルがあったとします。これをreturnで返す場合とyieldで返す場合（ジェネレータ関数）で比較してみましょう。

import csv

def read_csv_file_with_return(file_path):
    with open(file_path, 'r') as file:
        data = []
        for row in csv.reader(file):
            data.append(row)
        return data

def read_csv_file_with_yield(file_path):
    with open(file_path, 'r') as file:
        for row in csv.reader(file):
            yield row

returnを使った関数がList型変数dataへ行を追加して返すのに対し、yieldを使うと行ごとに返せます。
これはメモリ消費量の観点で非常に効率的です。

データベース接続を保持するとき

次にデータベース接続をするケースを見てみましょう。Pythonのデータベース系ライブラリのSQLAlchemyを使います。

import sqlalchemy as db

def get_db():
    engine = db.create_engine('postgresql://user:password@localhost/mydatabase')
    connection = engine.connect()
    try:
        yield connection
    finally:
        connection.close()