Pythonにおけるyieldやジェネレータについて調べてみた

先日バックエンドの開発をしていてyieldを使うと便利な場面があり、他にも使えるところないかなと思って調べたので、ついでにyieldの周りの知識や使い所についてまとめてみました。

どんな場面で使えるか

yield(というよりはジェネレータ)の使える状況を簡単に説明すると、主に次の1点だと思われます。

何らかの状態(ステート)を保持する必要があるが、クラスを作るまでもない、すぐにいらなくなるようなものを保持しないといけない時

具体的なケースでいうと

  • メモリを圧迫するぐらい大容量のcsvファイルを一行ずつ扱いたい時
  • 無限データストリームの生成
  • 再帰処理

といったことに使うことができます。


yieldとは

英語で「産み出す」や「もたらす」という意味があります。
yieldは、関数においてreturnの代わりにyieldを使用することによってその関数をジェネレータ関数にすることができます。
またasync def 関数で使用するとそのコルーチン関数は非同期ジェネレータになります。(今回非同期ジェネレータの話は割愛します)

# ジェネレータ関数
def generator_func():
    yield 1
    yield 2
    yield 3

# 非同期ジェネレータ
async def AsyncGenerator(num: int):
    for item in range(num):
        await asyncio.sleep(2)
        yield item

ジェネレータという言葉が出てきました。 ジェネレータの説明が必要ですね。

ジェネレータとは

Pythonにおける「ジェネレータ」とはジェネレータ関数を指す場合と、ジェネレータイテレータを指す場合があります。(言葉の定義があいまい。。)

ジェネレータ関数とは

ジェネレータ関数イテレータを返す、特殊な種類の関数のことです。
ジェネレータ関数によって作成されたイテレータのことを特にジェネレータイテレータと呼びます。
ジェネレータ関数はイテレータを簡単に作るためのものと言われています。(それだけではないですが)

# ジェネレータ関数
def generator_func():
    yield 1
    yield 2

# ジェネレータイテレータ
generator_iter = generator_func()
print(type(generator_iter)) # <class 'generator'>
print('__iter__' in dir(generator_iter)) # True
print('__next__' in dir(generator_iter)) # True

# 組み込みメソッドnext()もしくは特殊メソッド.__next__()で値を取得
print(next(generator_iter)) # 1
print(generator_iter.__next__())  # 2
# 要素を出し切ってからnextするとStopIterationエラーがraiseされる
print(next(generator_iter)) # raise StopIteration

# リストにも変換可能
# generator_iterとgenerator_iter_2の内部状況は違いに独立
generator_iter_2 = generator_func()
print(list(generator_iter_2))  # [1, 2]
# ジェネレータは一度使うとリストのように複数回使うことは不可
print(list(generator_iter_2)) # []

イテレータとは

イテレータ(オブジェクト)は「データの流れを表現」するオブジェクトで、値を 1 つずつ返してくれるオブジェクトのことです。
簡単には自身を戻り値とする__iter__()と次の要素を返す __next__() の 2 つのメソッドを持っています。
イテレータはリストやタプルなどのイテレータブルから作ることもできますし、クラスからも作ることができます。
イテレータイテレータブルの違いはこちら(Python公式Docs)

# イテレータブル(リスト)からのイテレータの生成
tmp_list = [1, 2, 3, 4]
print(type(tmp_list)) # <class 'list'>
new_iter = iter(tmp_list)
print(type(new_iter)) # <class 'list_iterator'>
print('__iter__' in dir(new_iter)) # True
print('__next__' in dir(new_iter))  # True
# クラスでのイテレータの定義
class PowTwo:
    def __init__(self, max=0):
        self.n = 0
        self.max = max

    def __iter__(self):
        return self

    def __next__(self):
        if self.n > self.max:
            raise StopIteration

        result = 2 ** self.n
        self.n += 1
        return result

pow_two_iter = PowTwo(2)
print('__iter__' in dir(pow_two_iter)) # True
print('__next__' in dir(pow_two_iter)) # True
print(next(pow_two_iter)) # 1
print(next(pow_two_iter)) # 2
print(next(pow_two_iter)) # 4
print(next(pow_two_iter)) # raise StopIteration

上と同様のものをジェネレータで実装するととても簡潔です。

def PowTwoGen(max=0):
    n = 0
    while n <= max:
        yield 2 ** n
        n += 1

pow_two_gen = PowTwoGen(2)
print(type(pow_two_gen)) # <class 'generator'>
print('__iter__' in dir(pow_two_gen)) # True
print('__next__' in dir(pow_two_gen)) # True
print(next(pow_two_gen)) # 1
print(next(pow_two_gen)) # 2
print(next(pow_two_gen)) # 4
print(next(pow_two_gen)) # raise StopIteration

ジェネレータイテレータとは(別名:ジェネレータオブジェクト)

ジェネレータ関数によって作成されたイテレータのことで、イテレータと同様__iter__()と次の要素を返す __next__() の 2 つのメソッドを持っています。
しかし、通常のイテレータとは異なる特徴として、yield文までの実行された、局所実行状態を保存して、処理を一時的に中断することができます。(ローカル変数、プログラムカウンタや評価スタック、そして例外処理のを含むすべてのローカルの状態が保持されます。)
そしてまたジェネレータイテレータが再開されると、中断した位置を取得し実行を再開します。
これはコルーチンとよく似た機構です。

このようにジェネレータは必要に応じて、計算を止めたり、再開することで、リストのように全ての要素を保持するのではなく、一つずつの要素をその都度、計算し生成するのでメモリを節約することができます。

その他の知識

ジェネレータについての簡単な説明は終わりますが、その他にもPythonのジェネレータには色々な機構があるのでそちらを紹介します。興味のない方は使用例まで飛ばしていただいて大丈夫です。

ジェネレータとコルーチン

ジェネレータ(イテレータ)は、中断と実行をしながら一時的にローカル状態を保存しつつ、値を返していくという性質を持っていて、それがコルーチンの機構と似ているということでした。
しかし、コルーチンはただ値を返すことだけではなくコルーチンに対して値を送信する機構があります。
Pythonのジェネレータではコルーチンと同様に値を送信することができ、これをジェネレータベースのコルーチンと言います(他の言語ではジェネレータにこの機構は存在しないのでこの特別な名称がついていると思われます。)

ジェネレータベースのコルーチン

ジェネレータに値を送信する場合は、next() 関数ではなくsend()メソッドを利用します。

※現在、ジェネレータベースのコルーチンはPython 3.10で廃止予定で、代わりにネイティブコルーチンの使用が推奨されています。
https://docs.python.org/ja/3/library/asyncio-task.html#generator-based-coroutines

def sumup():
    n = 0
    while True:
        n += yield n


coroutine = sumup()
# next() 関数を実行して、ジェネレータを実行状態に遷移
print(next(coroutine)) # 0

print(coroutine.send(1)) # 1
print(coroutine.send(2)) # 3
print(coroutine.send(3)) # 6

coroutine.close()

値を渡すsendのほかにもPythonのジェネレータはthrowとcloseというコルーチン特有のメソッドを使うことができます。

yield from

yield fromPython 3.3から登場した構文で、ジェネレータを作る時のfor文をyield fromを使うことでより簡潔に書くことができます。

# for文を使う場合
def generator_func():
    for i in range(1, 3):
        yield i

# yield forを使用(上と等価)
def generator_func_2():
    yield from range(1, 3)

しかしyield fromは上の例のようにfor文のシンタックスシュガーとして機能するだけではありません。
yield from を使うことでジェネレータの入れ子を簡単に実現することができ、ジェネレータオブジェクトをそのまま結合したようなふるまいをさせることができます。

def generator_A():
    for n in 'ABCDE':
        yield n

def generator_B():
    for n in '12345678':
        yield n

def generator_all():
    # 他のジェネレータ関数で作られたジェネレータオブジェクトが
    # yield した値がそのままyieldされる
    yield from generator_A()
    yield from generator_B()

print(list(generator_A())) # ['A', 'B', 'C', 'D', 'E']

print(list(generator_B())) # ['1', '2', '3', '4', '5', '6', '7', '8']

print(list(generator_all())) # ['A', 'B', 'C', 'D', 'E', '1', '2', '3', '4', '5', '6', '7', '8']

ジェネレータ式

ジェネレータ式とは簡単なジェネレータを作る時に使うことのできる書き方で、yieldを使用することなくジェネレータを作成できます。

# 普通のジェネレータの生成方法
def sample_generator():
    for x in range(5):
        yield x * 2
# ジェネレータ式を使った場合
generator_formula = (x * 2 for x in range(5))

print(next(generator_formula)) # 0
print(next(generator_formula)) # 2

ジェネレータ式はよくリスト内包表記と比較されます。ジェネレータ式はジェネレータイテレータを生成するのでリストとは違い、必要に応じてだけ値を算出するので、無限長ストリームや膨大なデータを返すようなイテレータを扱うことができます。

# https://docs.python.org/ja/3/howto/functional.html#generator-expressions-and-list-comprehensions

# 大規模データ
line_list = ['  line 1\n', 'line 2  \n', ...]

# ジェネレータ式
stripped_iter = (line.strip() for line in line_list)

# リスト内包表記
stripped_list = [line.strip() for line in line_list]

使い所

ここから本格的に使い方について紹介していきます。

大容量データの処理

すでにジェネレータ式で大容量のテキストデータを扱う例を書きましたが、よく使われるのでもう一度例をあげておきます。

# 大容量のファイルの場合以下はMemoryErrorが発生する場合がある。
def csv_reader(file_name):
    file = open(file_name)
    result = file.read().split("\n")
    return result

# ジェネレータ関数での書き換えを行い一行ずつ処理することができる
def csv_reader(file_name):
    for row in open(file_name, "r"):
        yield row

# ジェネレータ式の場合
csv_gen = (row for row in open(file_name))

無限のデータストリームの生成

あまり使うことはないかもしれませんが、無限長のデータストリームを生成することができます。 以下の例では理論的には全ての偶数を生成することができます。

def all_even():
    n = 0
    while True:
        yield n
        n += 2

処理のパイプライン化

複数のジェネレータを使って処理をパイプライン化することができます。よくわからないと思うので具体的な例を示します。

フィボナッチ数列を生成するジェネレータと数字を2乗するための別のジェネレータがあるとします。 フィボナッチ数列の二乗和を求めたい場合は、ジェネレータ関数の処理をパイプライン化することで、以下のように求めることができます。

def fibonacci_numbers(nums):
    x, y = 0, 1
    for _ in range(nums):
        x, y = y, x+y
        yield x

def square(nums):
    for num in nums:
        yield num**2

print(sum(square(fibonacci_numbers(15)))) # 602070

パイプライン化することによってきちんと処理が明示的に分割された可読性の高いコードが書くことができます。

再帰処理

ジェネレータはリスト全体を保持することなく、逐次処理ができるので、再帰処理にも利用されます。
以下の例はジェネレータを再帰的に使用し、木の走査を行う例です。

class Node:
    def __init__(self,val):
        self.val = val
        self.left = None
        self.right = None

# ジェネレータを使った木の走査
def inorder(node):
    if node:
        for x in inorder(node.left):
            yield x
        yield t.val
        for x in inorder(node.right):
            yield x

yield fromを使用するとより簡潔になります。

def inorder(node):
    if node:
        yield from inorder(node.left):
        yield t.val
        yield from inorder(node.right):

終わりに

コルーチンやら、非同期やらジェネレータとは関係なさそうな言葉が色々出てきました。言葉の定義の境目があいまいなものが多かったので、正確でない部分があるかもしれません。間違っていましたらフィードバックお願いします。 今回出てきた非同期ジェネレータなど非同期処理に関する内容が出てきたので、また自分の知識の整理のためにも非同期についての記事を書こうと思います。

参照一覧

https://www.lifewithpython.com/2015/11/python-create-iterator-protocol-class.html
https://docs.python.org/ja/3/howto/functional.html#generator-expressions-and-list-comprehensions
https://postd.cc/python-generators-coroutines-native-coroutines-and-async-await/
https://qiita.com/koshigoe/items/054383a89bd51d099f10
https://qiita.com/keitakurita/items/5a31b902db6adfa45a70#%E5%86%8D%E5%B8%B0
https://www.programiz.com/python-programming/generator