Python でジェネレータを作ったり、遅延評価してみる

Ruby の Enumerator でジェネレータを作ったり、遅延評価してみる」という記事を見かけたので Python 2 / 3 で書くとどのようになるか, 実際に書いてみることにしました.

イテレータ・ジェネレータの基本

Python におけるイテレータやジェネレータについてここで詳細は説明しないので「クロージャとジェネレータ – Dive Into Python 3 日本語版」や「クラスとイテレータ – Dive Into Python 3 日本語版」を参照してください.

from __future__ import print_function


def generator():
    yield 1
    yield 2
    yield 3


# Python 2 の場合
g = generator()
g.next()  # 1
g.next()  # 2
g.next()  # 3
g.next()  # Exception: StopIteration

# Python 3 の場合
g = generator()
next(g)  # 1
next(g)  # 2
next(g)  # 3
next(g)  # Exception: StopIteration


for i in generator():
    print(i)  # 順に 1, 2, 3 が表示されます

フィボナッチ数列のジェネレータ

例としてフィボナッチ数列のジェネレータを作成します. fib_generator はフィボナッチ数列を無限に返すジェネレータです.

from __future__ import print_function
from itertools import islice


def fib_generator():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b


# Python 2 の場合
fib = fib_generator()
fib.next()  # 0
fib.next()  # 1
fib.next()  # 1
fib.next()  # 2
fib.next()  # 3

# Python 3 の場合
fib = fib_generator()
next(fib)  # 0
next(fib)  # 1
next(fib)  # 1
next(fib)  # 2
next(fib)  # 3

islice(fib_generator(), 10)  # <itertools.islice at 0x...>
list(islice(fib_generator(), 10))  # [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

12-17行目はフィボナッチ数列を $latex F_0$ から $latex F_4$ まで順に取得する例です. 20-21行目はフィボナッチ数列の先頭から10個要素を取得する例です. islice はイタレータを返すので list 関数で list オブジェクトに変換する必要があります.

遅延評価

Python で yield を使ったジェネレータを作成すると遅延評価されていることはこれまでの例から分かることと思います. イテレータを引数にとってイテレータを返す関数を利用することで, より複雑な処理を行うことができます. このようなイテレータを生成する関数は標準ライブラリの itertools で多数提供されています.

次はフィボナッチ数列の要素をそれぞれ2乗した数列から, 奇数の要素だけの数列を作り, 先頭から10個取得する例です.

from itertools import islice


def fib_generator():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b


sq = lambda x: x ** 2
odd = lambda x: x % 2 == 1

islice(filter(odd, map(sq, fib_generator())), 10)  # 
list(islice(filter(odd, map(sq, fib_generator())), 10))  # [1, 1, 9, 25, 169, 441, 3025, 7921, 54289, 142129]

また Python 2 では map, filter などがイテレータではなくリストを返すため, 代わりにイテレータを返す itertools.imapitertools.ifilter を利用します.

from itertools import ifilter, imap, islice


def fib_generator():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b


sq = lambda x: x ** 2
odd = lambda x: x % 2 == 1

islice(ifilter(odd, imap(sq, fib_generator())), 10)  # 
list(islice(ifilter(odd, imap(sq, fib_generator())), 10))  # [1, 1, 9, 25, 169, 441, 3025, 7921, 54289, 142129]

まとめ

冒頭で示した Ruby での記事と比較して Python と Ruby の考え方の違いが垣間見えると思います.

Python 標準ライブラリの itertools のドキュメントには, 様々なイテレータを生成する関数や使用例がまとめられているので, 興味のある方は一度目を通すと良いです.

参考

更新

  • 2015/4/26 Python 2 / 3 での仕様の違いについて改善