【python】ジェネレータの作り方・使い方

ジェネレータとは

何もないところからイテレータを作り出すのがジェネレータです。
遅延実行されるのが大きな特徴です。
この特徴により、無限長の数列を表現できます。

そもそもイテレータとは

コレクション(リストやタプルなどの値の集合)から、1つずつ要素を取り出していくデザインパターンのことです。

イテレータを使った場合と使わない場合のfor文の具体例を見ていきます。

pythonのfor文は内部でイテレータを生成・使用しています。

lst = [0, 1, 2, 3]
# lstからイテレータを生成し、中身を1つずつ取り出している
for i in lst:
    print(i)

イテレータを使わないfor文はpythonでは書けないのでjsの例です。

const array = [0, 1, 2, 3]
const n = 4
for (let i = 0; i < array.length; i++) {
    // array全体に対して、インデックスでアクセスしている。
    console.log(array[i])
}

イテレータを使った場合は、コレクションからイテレータを通して、1つずつ値を取り出します。
一方、イテレータを使わない場合は、コレクション本体にインデックスでアクセスしたり、pop()などのメソッドを使用して値を取り出します。

作り方・使い方

簡単な例から初めて、最終的にフィボナッチ数列をジェネレータで実装していきます。

例1：1, 2, 3を返すジェネレータ

# ジェネレータ関数は普通の関数とは異なるのでそれが分かる命名がよい
def gen_123():
    print('egg')
    yield 1  # 1回目のnextでここまで実行
    print('and')
    yield 2  # 2回目のnextでここまで実行
    print('spam')
    yield 3  # 3回目のnextでここまで実行
    print('hum')  # 4回目のnextで実行されるが次にyieldがないのでStopIterationを投げる


if __name__ == "__main__":
    # ジェネレータ関数からジェネレータを生成
    gen = gen_123()
    # ジェネレータを1つ目のyieldまで実行
    first = next(gen)
    # firstには1回目のyieldの値が入る
    print(first)
    second = next(gen)
    print(second)
    third = next(gen)
    print(third)
    fourth = next(gen)  # !StopIteration

出力

egg
1
and
2
spam
3
hum
Traceback (most recent call last):
  File "hoge.py"
    fourth = next(gen)  # !StopIteration
StopIteration

next(gen)は、「次のyieldまで、genを進めてください」という命令です。
StopIterationは、「next(gen)されたのでジェネレータを進めたけど、次に返す値がありません」というメッセージです。

ちなみにですが、for文はinの右に書かれた式の結果からイテレータを生成し、そのイテレータに対してnext()をStopIterationが投げられるまで実行することで値を取り出しています。

例2:無限に1を返すジェネレータ

def gen_1():
    # 無限ループを作成
    while True:
        # next()されると常に1が返る
        yield 1
        # 無限ループなので、次のyieldが必ず存在する
        # つまり、gen_1()はStopIterationを返さない


if __name__ == "__main__":
    count = 0
    # gen_1()はStopIterationを返さないので無限ループになる
    for i in gen_1():
        # 無限ループ防止
        if count > 10:
            break
        print(i)
        count += 1

出力

1
1
1
1
・
・
・

無限ループとジェネレータ関数により、無限に続く数列を表現できました。

例3:フィボナッチ数列

import itertools


def gen_fib():
    """前の値(prev)と今の値(current)を足した値が次の値"""
    # 前の値
    prev = 1
    # 1回目のyieldでは1が取り出される
    yield prev
    # 今の値
    current = 1
    # 2回目のyieldも1が取り出される
    yield current
    # フィボナッチ数列には終わりがないので無限ループ
    while True:
        # 今の値を(今の値 + 前の値), 前の値を今の値に更新
        current, prev = current + prev, current
        # 更新された今の値が返る
        yield current


if __name__ == "__main__":
    # itertools.takewhileを使うことで、
    # 無限に続く数列から、条件から外れるまでの値を取り出せる
    # 100より小さいフィボナッチ数
    print(list(itertools.takewhile(lambda n: n < 100, gen_fib())))

    # itertools.isliceを使うことで、
    # 無限に続く数列からインデックスで指定した範囲を取り出せる
    # 101番目~110番目のフィボナッチ数
    print(list(itertools.islice(gen_fib(), 100, 110)))

    # itertools.isliceで生成したイテレータに対して、
    # 1回だけnextを適応することで、特定のインデックスの値を取り出せる
    # 100,001番目のフィボナッチ数は何桁？
    print(len(str(
        next(itertools.islice(gen_fib(), 100_000, None)))))

出力

[1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89]
[573147844013817084101, 927372692193078999176, 1500520536206896083277,
2427893228399975082453, 3928413764606871165730, 6356306993006846248183,
10284720757613717413913, 16641027750620563662096, 26925748508234281076009,
43566776258854844738105] ← 見やすいように改行している
20899

フィボナッチ数列は無限に続く数列です。なのでlist(gen_fib())と書くと、無限ループに陥ってしまいます。
今回の例では、イテレータ構築用の標準ライブラリitertoolsを使って、新たに有限なイテレータを生成しリストにすることで、実際に値を取り出しています。

最後に

ジェネレータが活躍する場面として、無限に続くフィボナッチ数列から特定の値を取り出すような処理を紹介しました。
これを一般化して言うと、ジェネレータが活躍する場面はリスト(メモリ)に載せることが難しいor不可能な値のシークエンスから、その一部を取り出して処理する必要がある場合、になるかと思います。

この記事がpythonでプログラミングするときに、ジェネレータのことを思い出す一助になればと思います。