Pythonのイテレータ(iterator)とジェネレータ(generator)についてまとめてみます。

(追記2018.12.25: Python3の文法に全面的に置き換えました)

イテレータ: 要素を反復して取り出すことのできるインタフェース
ジェネレータ: イテレータの一種であり、1要素を取り出そうとする度に処理を行い、要素をジェネレートするタイプのもの。Pythonではyield文を使った実装を指すことが多いと思われる

Python組み込みのコレクション(list, tuple, set, dictなど)はどれもイテレーション可能ですが、組み込みのコレクションを使った繰り返し処理ではあらかじめコレクションに値を入れておく必要があるため、以下のようなケースではイテレータやジェネレータを自分で実装したいというケースがあると思います。

無限に繰り返すイテレーション
要素すべてをあらかじめ計算しておく/取得してくるのが計算コスト/処理時間/メモリ使用量などの面で大変

クラスによるイテレータの実装

あるオブジェクトをforのinなどのイテレータを期待するコンテクストに置くと、まずオブジェクトの__iter__()メソッドが呼ばれ、イテレータ実装を返すことが求められます。この返り値で得られたオブジェクトは__next__()というメソッドが呼ばれます。__next__()はStopIteration例外が出るまで呼ばれます。

普通にlistとかわりませんが、インスタンス生成時に与えられた数字のリストを順番に返す実装の例です。

sample1.py

class MyIterator(object):
	def __init__(self, *numbers):
		self._numbers = numbers
		self._i = 0
	def __iter__(self):
		# __next__()はselfが実装してるのでそのままselfを返す
		return self
	def __next__(self):  # Python2だと next(self) で定義
		if self._i == len(self._numbers):
			raise StopIteration()
		value = self._numbers[self._i]
		self._i += 1
		return value
        
my_iterator = MyIterator(10, 20, 30)
for num in my_iterator:
	print('hello %d' % num)

結果は

hello 10
hello 20
hello 30

になります。

この例では__iter__()はselfを返してしまっていますが、イテレーションのための処理が複雑になりそうなときなどは別にイテレーションのための実装クラスを実装して、そういったオブジェクトを生成して返すというのもありということですね。

組み込み関数iter()を使うとlistなどの組み込みの型もこのルールに従って実装されていることがわかります。

>>> hoge = [1, 2, 3]
>>> hoge_iter = iter(hoge)
>>> hoge_iter.__next__()
1
>>> hoge_iter.__next__()
2
>>> hoge_iter.__next__()
3
>>> hoge_iter.__next__()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

イテレータまとめ

オブジェクトに対してイテレータ化が要求されたとき __iter__() メソッドが呼ばれる
__next__()メソッドは呼ばれる度に新しい値を返す
__next__()メソッドはもう返す値がないときの呼び出しではStopIteration例外をraiseする

yieldを使ったジェネレータの実装

yieldは慣れないと理解しずらいかもしれませんが、仕組みは簡単です。yieldを使ったジェネレータの実装ではわざわざクラスを定義する必要がありません。次のようなジェネレータ関数を定義してみます。

my_generator.py

 def my_generator():
	yield 1
	yield 2
	yield 3

これは1、2、3という3つの値を順番に生成してくれるジェネレータです。

ジェネレータは以下のような計算コストが問題となるケースで使われることが多いです。

あらかじめ繰り返す値をすべて計算しておけない
計算コストを節約するために要素ごとに生成の計算を行う

ジェネレータ関数は関数呼び出しを行うことでイテレータオブジェクトになります。

gen = my_generator()
gen.__next__()  # 1
gen.__next__()  # 2
gen.__next__()  # 3
gen.__next__()  # StopIteration

yieldすると、next()を呼んだ側に制御が戻ります。以下のようにprint文をはさんで処理の流れを確認してみます。

generator_sample.py


def my_generator():
	print('before yield')
	yield 1
	print('yielded 1')
	yield 2
	print('yielded 2')
	yield 3
	print('yielded 3, finished')

def main():
	gen = my_generator()
	print('start')
	v1 = gen.__next__()
	print('called __next__(), v1=%s' % v1)
	v2 = gen.__next__()
	print('called __next__(), v2=%s' % v2)
	v3 = gen.__next__()
	print('called __next__(), v3=%s' % v3)
	v4 = gen.__next__()  # should be exception

main()

実行結果は以下のようになります。

start
before yield
called __next__(), v1=1
yielded 1
called __next__(), v2=2
yielded 2
called __next__(), v3=3
yielded 3, finished
Traceback (most recent call last):
  File "./generator_sample.py", line 21, in <module>
    main()
  File "./generator_sample.py", line 19, in main
    v4 = gen.__next__()  # should be exception
StopIteration

ジェネレータまとめ

ジェネレータの実装ではyieldを使う
yieldした回数だけ値が出てくる
ジェネレータ関数を関数呼び出しするとイテレータオブジェクトになる

list/tuple/set/リスト内包表記とイテレータ

イテレーション可能なものはlist/tuple/set/リスト内包表記などの組み込みの機能と簡単に連携させることができます。

たとえば、上記の簡単なジェネレータをlist()関数に渡すと簡単に[1, 2, 3]といった値のlistオブジェクトに変換できます。tupleやsetも同様です。


def my_generator():
	yield 1
	yield 2
	yield 3

def my_generator2():
	yield 10
	yield 10
	yield 20

print(list(my_generator()))  # => [1, 2, 3]
print([v * 2 for v in my_generator()])  # => [2, 4, 6]
print(set(my_generator2()))  # => set([10, 20])

もちろん yield で実装されたジェネレータだけでなく __next__() と return で実装したイテレータも同様に組み込みの機能と連携させることができます。

itertools

Pythonにはitertoolsというイテレータオブジェクトを組み合わせて様々な操作を簡単に行うことのできるライブラリがあるのでご紹介しておきます。これは自分でイテレータを実装しなくてもlist/tupleなどの組み込みデータでの利用がメインと思いますので、イテレータを実装していなくても便利です。

たとえば[1, 2, 3]と['a', 'b']のすべての組み合わせを列挙させることなどが簡単にできたりします。

more_itertools

more_itertoolsという標準ではないPyPIライブラリもあります。itertoolsに入っていない、便利な関数が多数収録されています。N個ごとに塊にする chunked やイテレータを回して個数をカウントする ilen などがあります。

利用するにはインストールする必要があります。

$ pip install more-itertools

itertools/more_itertoolsを紹介・解説してくださっているQiita記事もありますので参考になるかと思います。

ジェネレータを何回も繰り返し使う

ジェネレータは一度forループで回したりすると2回目以降のforループでは要素が出てきません。

ジェネレータ関数の呼び出しに副作用がなく何回でも呼び出せるようにしたい場合は「
Python のジェネレータを何回もイテレートしたい」に記述したテクニックを使うと便利かもしれません。

まとめ

Pythonにおけるイテレータインタフェースは__next__()であり, StopIteration例外でもう要素がないことがわかる
オブジェクトはイテレータコンテキストとして評価されると__iter__()が呼ばれる
ジェネレータはイテレータの一種であり、 yield を使って実装する
ジェネレータ関数は呼び出すとイテレータオブジェクトになる
無限に繰り返したりあらかじめ全部計算しておけない場合にイテレータを実装すると便利なことがある
list/tuple/set/リスト内包表記などと組み合わせるとさらに便利
itertools もとても便利

参考リンク

Iterator Types: https://docs.python.org/3/library/stdtypes.html#typeiter
itertools: https://docs.python.org/ja/3/library/itertools.html
more_itertools: https://more-itertools.readthedocs.io/en/stable/