2015/11/24

Python のイテレータ生成クラスの使い方

今回は Python のイテレータ生成クラスを使う方法についてご紹介します。

Python でイテレータを生成する関数のことを「ジェネレータ」と呼びますが、関数と同様にクラスを使う形でもイテレータを生成することができます。

具体的には __iter__() メソッドを定義してこれがイテレータを返すようにすれば OK です。
yield を使わない場合は __iter__() とは別に __next__() ( Python 2 の場合は next() )メソッドを定義する必要がありますが、シンプルなケースでは __iter__() と yield だけで十分かと思います。

与えられた整数の因数を返すジェネレータクラスを書いてみます。

# coding: utf-8

# 因数を返すイテレータクラス
class PrimeFactor(object):
   
    // number で対象の数を、 max_count で返す因数の数を指定する
    def __init__(self, number, max_count):
        self.number = number
        self.max_count = max_count

    // クラスをジェネレータ化する
    def __iter__(self):
        count = 0
        for i in range(2, self.number):
            if self.number % i == 0:
                count += 1
                // 指定された数だけ因数が返されたらイテレータ処理を終了するために StopIteration() 例外をあげる
                if count > self.max_count:
                    raise StopIteration()
                yield i

ここで定義した PrimeFactor は以下の形で使うことができます。

pf = PrimeFactor(number = 100, max_count = 5)

# 1 周目
for n in pf:
    print(n, end=", ")
print()
# => 2, 4, 5, 10, 20,

# 2 周目
for n in pf:
    print(n, end=", ")
print()
# => 2, 4, 5, 10, 20,

実装方法の問題でもあるかと思いますが、イテレータを返す関数とのちがいとして、クラスの場合はいったんループが回りきった後にもう一度ループを開始できるという特徴があります。

// PrimeFactor の関数版
def prime_factor(number, max_count):
    count = 0
    for i in range(2, number):
        if number % i == 0:
            count += 1
            if count > max_count:
                raise StopIteration()
            yield i

// こちらはいったんループを回したらその後は StopIteration 例外を出し続けるため
// 2 周目のループは実行されない
pf2 = prime_factor(number = 100, max_count = 5)

# 1 周目
for n in pf2:
    print(n, end=", ")
print()
# => 2, 4, 5, 10, 20,

# 2 周目
for n in pf2:
    print(n, end=", ")
print()
# => 出力なし

ちなみにこのイテレータを生成するクラスが持つインタフェースのことを「イテレータプロトコル( iterator protocol )」と呼んだりもするそうです。

以上です。

参考:

- 3.5 イテレータ型
- python - How to make class iterable? - Stack Overflow

0 件のコメント: