実行時の環境

Python
3.11.5

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

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

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

yieldとreturnの違い

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

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

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

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

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

python
def sample():
    yield 1
    yield 2
    yield 3

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

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

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

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

python
def sample():
    yield 1
    yield 2
    yield 3

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

List（配列）型へ変換

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

python
def sample():
    yield 1
    yield 2
    yield 3

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

結果は以下です。

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

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

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

python
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))

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

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

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

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

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

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

python
def sample():
    yield 1
    yield 2
    yield 3

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

結果は以下です。

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

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

python
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例外が発生します。

python
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))

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

shell
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型だと[]ですが、ジェネレータの場合は()で囲みます。

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

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

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

shell
1
4
9
16
25

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

ジェネレータの使い所

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

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

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

python
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を使います。

python
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()

データベース接続でreturnではなくyieldを使う理由は、CSVファイルと同じく多大なデータ量になった際、メモリを効率的に使うためです。
データベース接続によって得られるデータは、たいていの場合はイテラブル（繰り返し処理できるオブジェクト）なのでジェネレータが有効となります。

その他の例

他にもジェネレータの利用が有効なシーンはたくさんあります。コードを書くと長くなるため、以下に列挙します。

  • Webサーバへのリクエストのような大量の順次処理が発生するケース
  • 複雑な計算処理を並行で行うケース
  • 形態素解析
  • データパイプライン処理
  • 再帰処理

ジェネレータを覚えて効率的な実装をしよう

なかなか使うシーンを想像しづらいジェネレータですが、多くのWebアプリケーションフレームワークやライブラリで処理の効率化を考慮して利用されています。
自身で実装するアプリケーションやツールでも、効率的に利用できるケースがあったら積極的に使っていきましょう。

メモリ消費を抑える実装は、アルゴリズムや設計といった基礎的なスキル以外にも、こういった言語の機能でサポートされていることがあるので、ぜひ覚えておきたいところです。