Python で iterator を使うときに考える初期化のタイミング

Python のクラスで __iter__ や __next__ を定義すると、そのインスタンスは for 文などで利用できます。

for 文で使うインスタンスを作るクラスを作ってみます。next を呼び出した回数をカウントしながら、1ずつ増やした数値を返し、規定の回数に到達したら StopIteration を出してループを終了させるようにします。

forで使うクラス

class Iterator:

    def __init__(self):
        print("init")
        self.chunk = 3
        self.number = 0

    def __iter__(self):
        print("iter")
        self.count = 0
        return self

    def __next__(self):
        if self.count >= 3:
            raise StopIteration()

        self.count += 1
        self.number += 1
        return self.number

インスタンスは、複数回にわたって for 文に渡すことができます。

__init__ が呼ばれるときはインスタンスを定義した最初のみで、 __iter__ は for 文に使う度に呼ばれます。出力を見ると、 print("init") が最初に、 __iter__ が first, second, third の出力の前にそれぞれ出現しているのがわかります。

forを3回用いる

iterator = Iterator()

for i in iterator:
    print("first:", i)

for i in iterator:
    print("second:", i)

for i in iterator:
    print("third:", i)

出力

init
iter
first: 1
first: 2
first: 3
iter
second: 4
second: 5
second: 6
iter
third: 7
third: 8
third: 9

ここで、クラス内の self.count は __iter__ が呼ばれているため3回リセットされていますが、 self.number は __init__ で初期化しているだけのため、前回のループの値から引き続いて数値が増えていっています。

毎回の for で、 1,2,3 を出力したい場合、今回のように単純なクラスを用いる場合は、毎回インスタンスを作るのでも良いかもしれません。

毎回インスタンスを作る

for i in Iterator():
    print("first:", i)

for i in Iterator():
    print("second:", i)

for i in Iterator():
    print("third:", i)

ただし、インスタンスの中でファイルの入出力をしたり、複雑な計算の結果を保存したりしている場合、そのインスタンスの使い回しをしたい場合もあるでしょう。

その場合は、 __iter__ に毎回の初期化処理をしっかり記述することが必要です。

具体例として、 gensim の Word2Vec に文章を渡すとき、容量が多いためファイルを読みながら形態素解析をして、文章を一つずつ書き出していくときに使えます。このとき、イテレーターのインスタンスを引数に与えるわけですが、 Word2Vec は学習時に複数回ループを回すため、 __init__ に全ての初期化処理を記述してしまうと、2回目以降の学習データが全部空になってしまいます。

Python の記事を調べていると、多くのコードで def __iter__(self): の中身が return self だけになっているため、筆者もあまり意識していませんでした。場合によっては致命的な不具合に繋がりかねないため、注意が必要です。