GeneratorExitを捕捉してジェネレータ側からイテレーションの終了を検知する

イテレータそのものはbreak等で最後まで回らない可能性があるけど、ジェネレータ内部のイテレーションは最後まで回したいという場合がまれにある。

うまくいかない例

例として次のrange_printsumはrange関数にその範囲の総和をprintする処理を追加したものだ。

gen.py

def range_printsum(*args, **kwargs):
    sum_ = 0

    for i in range(*args, **kwargs):
        sum_ += i
        yield i

    print(sum_)

このジェネレータ実装ではイテレーションを途中で止めてしまうと、sum_に格納されるのは実際にyieldした数の合計だし、そもそもsum_は出力されない。

>>> from gen import range_printsum
>>> for i in range_printsum(11):
...     pass
...
55
>>> for i in range_printsum(11):
...     if i > 5:
...         break
...
>>>

うまくいく例

そこでGeneratorExitを捕捉する。これはジェネレータが閉じられた（ジェネレータのclose()が呼ばれた）時に送出される。これはfor等のループが中断された場合でも捕捉できる。
GeneratorExitを捕捉した後は値を返さず、加算だけ行うようrange_printsumを修正する。

gen.py

def range_printsum(*args, **kwargs):
    sum_ = 0
    stop_yield = False

    for i in range(*args, **kwargs):
        sum_ += i
        if not stop_yield:
            try:
                yield i
            except GeneratorExit:
                stop_yield = True

    print(sum_)

今度は期待通りに動く。

>>> from gen import range_printsum
>>> for i in range_printsum(11):
...     pass
...
55
>>> for i in range_printsum(11):
...     if i > 5:
...         break
...
55
>>>

このように中断しても動く。

>>> try:
...     for i in range_printsum(11):
...         if i > 5:
...             raise IndexError
... except IndexError:
...     pass
...
55
>>>

ジェネレータの仕様による注意点として、一度閉じられたジェネレータは値を生成してはならない。GeneratorExitを捕捉した後に何らかの値をyieldするとRuntimeError: generator ignored GeneratorExitというエラーが発生する。