Python で iter と next を実装する

Python でオブジェクトを for 文で使えるようにするには、「イテレータプロトコル」を実装します。具体的には __iter__() と __next__() という2つの特殊メソッドを定義します。

イテレータプロトコルとは

イテレータプロトコルは以下の2つのメソッドで構成されます。

__iter__()

イテレータオブジェクト自身を返す。iter() 関数から呼び出される。

__next__()

次の要素を返す。要素がなければ StopIteration を送出する。next() 関数から呼び出される。

基本的な実装

1から指定した数までカウントするイテレータを作ってみましょう。

class Counter:
    def __init__(self, max_value):
        self.max_value = max_value
        self.current = 0
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current >= self.max_value:
            raise StopIteration
        self.current += 1
        return self.current

counter = Counter(3)
for num in counter:
    print(num)
# 1
# 2
# 3

動作の仕組み

for 文を使うと、内部で以下の処理が行われます。

iter() が iter() を呼ぶ

next() が next() を呼ぶ

StopIteration で終了

手動で操作することもできます。

counter = Counter(2)
it = iter(counter)     # __iter__() が呼ばれる

print(next(it))        # __next__() → 1
print(next(it))        # __next__() → 2
print(next(it))        # __next__() → StopIteration

イテラブルとイテレータの分離

上の例ではオブジェクト自身がイテレータですが、イテラブルとイテレータを分離する設計もあります。これにより、同じオブジェクトを何度でもイテレートできます。

class CounterIterable:
    def __init__(self, max_value):
        self.max_value = max_value
    
    def __iter__(self):
        return CounterIterator(self.max_value)

class CounterIterator:
    def __init__(self, max_value):
        self.max_value = max_value
        self.current = 0
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current >= self.max_value:
            raise StopIteration
        self.current += 1
        return self.current

counter = CounterIterable(3)

# 何度でもイテレートできる
print(list(counter))  # [1, 2, 3]
print(list(counter))  # [1, 2, 3]

ジェネレータを使った簡略化

実際には、__iter__() でジェネレータを返す方が簡潔に書けます。

class Counter:
    def __init__(self, max_value):
        self.max_value = max_value
    
    def __iter__(self):
        for i in range(1, self.max_value + 1):
            yield i

counter = Counter(3)
print(list(counter))  # [1, 2, 3]