[Python] イテレータを実装する #Python3

イテレータとは

まず始めにイテレータとは何かを、実際に動かしてみて確認しましょう。
リストからイテレータを使用して要素を一つずつ取り出してみます。

>>> l = [0, 1, 2]
>>> it = iter(l)

リストlはイテレータではないので、要素を順番に取り出すために組み込み関数iter()を使用してイテレータオブジェクトを取得します。

>>> next(it)
0
>>> next(it)
1
>>> next(it)
2

イテレータオブジェクトは組み込み関数next()を使用すると要素を一つずつ取り出すことができます。

>>> next(it)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>>

そして、取り出す要素がなくなった場合にはStopIteration例外が発生します。
このようにしてリストから要素を一つずつ取り出すことができます。これがイテレータの使用方法になります。

しかし実際にはiter()やnext()は使用せず、以下の様にfor文を使用することが多いと思います。

>>> l = [0, 1, 2]
>>> for i in l:
...     print(i)
...
0
1
2

この場合、for文の内部でリストlからiter()によりイテレータを作成し、StopIterationが発生するまでnext()を呼び出し要素を一つずつとりだしています。そしてStopIterationが発生したらループ処理を終了します。

イテレータを実装する

それではイテレータの機能を持つクラスを実装してみます。クラスにイテレータの機能を持たせるには、__next__()メソッドを持つオブジェクトを返却する__iter__()メソッドを定義すればイテレータとして動作します。
以下のクラスは0から指定した数までインクリメントした数を返却するクラスになります。

sample_iterator.py

class SampleIterator(object):
    def __init__(self, num):
        self.num = num
        self.current = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.current == self.num:
            raise StopIteration()

        ret = self.current
        self.current += 1
        return ret

実際に使用してみます。最初はnext()を使用して呼び出します。

>>> from sample_iterator import SampleIterator
>>> si = SampleIterator(3)
>>> next(si)
0
>>> next(si)
1
>>> next(si)
2
>>> next(si)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
    raise StopIteration()
StopIteration

最初に行ったiter()を使用したものと同様の動作になりました。組み込み関数のnext()で呼ばれたときに__next__()メソッドが呼ばれていることがわかります。
次にfor文で使用してみます。

>>> from sample_iterator import SampleIterator
>>> si = SampleIterator(3)
>>> for i in si:
...     print(i)
...
0
1
2

__iter__()メソッドが最初に呼ばれ、その後__next__()メソッドが呼ばれています。

まとめ

イテレータの機能を持つクラスを実装するには

__next__()メソッドを持つ__iter__()メソッドを実装する
__next__()メソッドは返却する要素がなくなったらStopIterationを発生させる

ソースコードsample_iterator.py