Pythonのgenerator使用時の注意点 (generatorの再利用)

こんにちはtatsyです。

タイトルの通り、Pythonでgeneratorを使うときの注意点で気づいたことがあったので備忘録的にまとめておきます。

もしgeneratorが何かは知っていて、結果だけご覧になりたい場合には、お手数ですが下までスクロールをお願いいたします。

Generatorとは?


Pythonにおけるgeneratorとは、リストのような変数の並びを順番にたどるための対象の一種です。

例えば、Pythonであれば、

lst = [ 1, 2, 3, 4 ]
for item in lst:
  print(item)

のようにしてリストの内容をたどることができます。ですが、リストの場合にはたどりたい対象が変数として計算済みである必要があります。

例えば、長さが10000の32bit整数の配列をたどろうとするならば、10000 x 32 bit = 320000 bit = 40000 byte = 40KBのメモリが必要になります。

基本的にメモリ領域の確保には時間がかかりますので、例えばfor文で1-10000の数をたどるたびに上記のリストを確保していては時間がかかりすぎるわけです。

実はPython 2.x系のrange関数は内部でリストを作るようになっており、これがfor文実行が遅い原因の一つでした(その代わりにxrangeというのがありましたが)。

ちなみにPython 2.7でrange(1000000000)のような値を評価しようとするとコンソールが帰ってこなくなると思います。

一方でgeneratorは値を必要になったときに確保します。上記の1-10000の数列の例で言えば、必要なメモリ領域は当然32bit整数1つ分で32bitです。もう少し専門的に言えば、数列の値が遅延評価されており、その分だけメモリ消費を抑えられます。

Python 3.x系のrange関数はこのgeneratorを採用しておりまして、どんなに長い数をrangeに入力しても基本的には大丈夫です。

Python 2.7での場合と違いPython 3.4でrange(1000000000)を実行しても、一瞬でコンソールが帰ってきます。これがリストとgeneratorの違いです。

geratorの作り方


Generatorの一般的な作り方は関数の中でyieldを呼び出す方法です。例えば1から100までの値を返すようなgeneratorを作りたければ、

def iter100():
  i = 1
  while i <= 100:
    yield i
    i += 1

のような関数を作れば良いです。使い方は普通のリストの場合と同じで、

for i in iter100():
  print(i)

のようになります。

ちなみにリスト内包表記でもgeneratorを作ることができます。通常のリストを生成する場合には四角いカッコ[]で評価する式を囲いますが、generatorを生成する場合には丸いカッコ()で評価する式を囲います。

lst = [ i for i in range(100) ]  # 0-99のリスト
gen = ( i for i in range(100) )  # 0-99を戻すgenerator

itertoolsパッケージ


Generatorを語る上で外せないのがitertoolsパッケージです。Python 3.5.2の公式ドキュメントはこちらになります。

http://docs.python.jp/3/library/itertools.html

itertoolsパッケージはイテレーションをする上で便利なgeneratorを生成するための関数がまとめられたパッケージになります。代表的なものをいくつか見ていきます。

count

count関数は引数として1つの整数を取って、その整数をスタートとして1ずつ増えていくような無限の長さの数列を生成します。

from itertools import count

for i in count(-5):
  print(i)
  if i >= 5: break

# -5, -4, ..., 4, 5が出力される

こちらのcount関数はgeneratorの特徴である遅延評価の特性をよく表しています。通常のリストのように先に評価が必要な方法(正格評価)では無限の長さのたどるような対象を作ることはできません(メモリが有限なので)。

一方でgeneratorを使う場合には、必要なのは今の数がいくつであるのかと、次の数を計算するための方法なので、結果的に長さが無限になるとしても基本的にはOKです。

ただし、上記の例のように何処かの数字でbreakをしてあげないとfor文が無限に回り続けることになりますので注意が必要です。

ちなみにcountを自分で実装するとすると、以下のようになるかと思います。

def count(n):
  while True:
    yield n
    n += 1

product

productは画像処理のように二重のforループを使う場合には必須ではないかなと個人的には思っております。

productは2つのリスト、あるいはgeneratorの要素のすべての組み合わせをたどるようなgeneratorを生成します。次の例をご覧ください。

for i, j in product(range(3), range(2)):
  print(i, j)

# 0, 0 -> 0, 1 -> 1, 0 -> 1, 1 -> 2, 0 -> 2, 1の順で出力される

こちらのほうがfor文のネストが少なくコードの見通しが良くなりますが、ideoneでコードを書いてみると速度はあまり変わりませんでした(今までproductのほうが速いのかと思っていました)。

http://ideone.com/mVksYn

こちらもPythonのコードで実装したとすると以下のようになります。

def product(gen1, gen2):
  lst1 = [ i for i in gen1 ]
  lst2 = [ j for j in gen2 ]
  for i in lst1:
    for j in lst2:
      yield i, j

こちらのコードでは最初にリスト内包表記を使って内容を評価していますが、実際にproduct関数の引数にrange(1000000000)のような大きな値を与えるとコンソールが帰ってこなくなることからも、内部で一度評価が行われていることがわかるかと思います。

実は、ここで評価を行わないと行けない理由こそが、今回ご紹介したい注意点につながっていたりします。

Generatorを使うときの注意点


まず、当たり前のこととして、何度も同じ数列をたどるのであればリストの値を持っておいたほうが早くなります。generatorの場合には、使うたびに評価が入るからです。

では、generatorを繰り返し使うことはそもそも可能なのか?というと答えはNoです。

試しに次のコードを実行してみてください。

gen = ( i for i in range(10) )
for i in gen:
  print(i)

for i in gen:
  print(i)

そうすると0から9の数字が1回しか表示されないのがわかるかと思います。

これはなぜかというと、generatorは内部の変数の状態を保存しているためです。上記のgeneratorが一度動き終わるとiの値はすでに10になっていて、終わりの数である10以上になってしまいます。

つまり、そもそも遅くなるので使わないほうがいいし、使うこともできない(使えるが動作がおかしくなる)ということなのです。

これに最近まで気づいていなくて、バグをだしてしまいました…。内部の動きを考えればある程度当たり前ではあったと思うんですけどね。

ですので、自分でgeneratorを自作する場面で勘違いを防ぐのであれば、

def range(n):
  i = 0
  while i < n:
    yield i
    i += 1
  raise StopIteration()

のように例外を投げておく方がいいかもしれません。

あるいは、通常のgeneratorの使い方を変更して、

gen = ( i for i in range(10) )
while True:
  try:
    i = next(gen)
    print(i)
  except StopIteration:
    break

のようにnextで次の要素を取るようにすれば、次の要素がない時にはStopIteration例外が投げられますので、勘違いを防ぐことができます。

こっちは書き方が少し複雑になるので、コードをきれいにしておくのであれば個人的には前者かなと思います。

まとめ


というわけでgeneratorの概要および最近発見した使用時の注意点についてご紹介しました。何かのお役に立てば幸いです。

今回も最後までお読みいただきありがとうございました。