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]
この方法なら __next__() を自分で実装する必要がありません。