pythonでyield先で起きた例外をyield元で取り扱いたい

「pythonでyield先で起きた例外をyield元で取り扱いたい」というようなことを言っていた。いろいろ調べ直してまとめてみたらgenerator周りの復習になったのでそのメモ。

課題

yeild先で起きた例外をyield元で取り扱いたい？とはなにか？

def _use(i):
    if i % 2 == 0:
        raise Exception("oops")
    print(i)

# yield元
def gen():
    yield 1
    yield 2
    yield 3

for i in gen():
    _use(i) # yield先

ここでは、generatorが値を生成している部分をyield元、生成された値を消費している部分をyield先と表現している。

当たり前だけど、通常のgeneratorの利用では、例外の発生地点でハンドリングすることになる(つまりyield先でハンドリングするということ)。場合によっては代わりにyield元でハンドリングしたくなることがある。

コレをどうにかできないか？というのが今回の課題。

generator.throw()

generatorにはsend()のほかにthrow()が存在する¹。コレを使えばyield元に委譲することができそう。

すごく雑なgeneratorのprotocol。このthrow()が今回の主題。

class Generator(Protocol):
    def send(self, value:Any) -> Any:
        ...

    def throw(self, typ:Type[Exception], value:Optional[Exception]=None, traceback:Optional[Traceback]=None) -> Any:
        ...

throw()の説明を読んでみると以下の様なことが書かれている。

ジェネレータが中断した位置で type 型の例外を発生させて、そのジェネレータ関数が生成する次の値を返します。ジェネレータが値を生成することなく終了すると StopIteration が発生します。ジェネレータ関数が渡された例外を捕捉しない、もしくは違う例外を発生させるなら、その例外は呼び出し元へ伝搬されます。

ちょっと文意が読み取りづらいが、以下の文章が挙動の把握として重要な点。

ジェネレータが値を生成することなく終了すると StopIteration が発生します

いろいろ実験してみよう。

throwされた例外のexcept部分で、yield元が値を生成しない場合

generatorは1,2,とやってきてその次はNoneを返す。yield先では数値を期待して加算しようとするのでTypeErrorが発生する。それをthrow()で制御を返してあげている。そしてyield元では、発生した例外に対してprintするだけ。

このようなコードの場合には、例外への対応が終わった段階でStopIterationが発生して。中断される。

def gen():
    yield 1
    yield 2
    try:
        yield None
    except Exception as e:
        print("hmm", e)
    yield 4  # never


def use(gen):
    def _use(n):
        print(n + 1)

    for i in gen:
        try:
            _use(i)
        except TypeError as e:
            gen.throw(e)


use(gen())

実行結果は以下の通り。4はyieldされない。

2
3
hmm unsupported operand type(s) for +: 'NoneType' and 'int'

throwされた例外のexcept部分で、yield元が値を生成する場合

今度は少しコードを変えてみよう。throwされた例外を補足したexceptのところで、yieldしてみる。100をyieldしてみよう。今度は、中断せずその先まで呼ばれる。つまり5もyieldされる。

def gen():
    yield 1
    yield 2
    try:
        yield None
    except TypeError as e:
        print("hmm", e)
        yield 100
    yield 5  # ok


def use(gen):
    def _use(n):
        print(n + 1)

    for i in gen:
        try:
            _use(i)
        except TypeError as e:
            print("!", gen.throw(e))


use(gen())

実行結果は以下の通り。

2
3
hmm unsupported operand type(s) for +: 'NoneType' and 'int'
! 100
6

これで、ドキュメント中の「ジェネレータが値を生成することなく終了すると StopIterationを送信します」ということの意味がわかったと思う。つまり、yield元で例外を補足した際のexceptの部分でyieldするかしないかで中断するかしないかが決まる。

yield先でのbreakをyield補足するには？

もうちょっと問題を変えてみよう。yield先での例外ではなく、yield先がbreakしたことを、yield元が知るにはどうしたら良いんだろうか？コレも実はドキュメントに答えが書いてある。close()を参照。

ジェネレータ関数が一時停止した時点で GeneratorExit を発生させます。そして、ジェネレータ関数が無事に終了するか、既にクローズされているか、(例外が捕捉されなかったために) GeneratorExit が送出された場合、 close は呼び出し元へ戻ります。ジェネレータが値を生成する場合 RuntimeError が発生します。 close() はジェネレータが例外や正常な終了により既に終了している場合は何もしません。

今回は、GeneratorExitという例外について把握すれば良い。

そんなわけでこの例外を補足するコードを書いてみる。1,2,3とyieldされていくが、yield先で3が消費されたタイミングでbreakが呼ばれる。4,5はyieldされない。

def gen():
    try:
        yield 1
        yield 2
        yield 3
        yield 4  # never
        yield 5  # never
    except GeneratorExit as e:
        print("!?", e)


def use(gen):
    def _use(n):
        print(n + 1)

    for i in gen:
        _use(i)
        if i == 3:
            break


use(gen())

実行結果は以下。

2
3
4
!?

しっかりとbreakをyield元が補足できている。

中断を検知できる嬉しさ

breakを補足できたところで何が嬉しいの？と思うかもしれない。でもこれはけっこう大切なことで、例えばcoroutine²としてgeneratorを使った場合に、途中で中断されたあとに、後始末をしたくなる事がある。例えば、開いていたコネクションを閉じたりなど。コレができると本当にいろいろできる。

gist

https://gist.github.com/podhmo/f6e2c4dc2d645a7c4411c525b8722263

throwの存在は、methaneさんに教えてもらった。 https://mobile.twitter.com/methane/status/1325302752944160768 ↩
ちなみにgeneratorのこの辺の仕様のPEPには、本来の意味でのcorountineの用法(?)でcorountineの言葉が出てくる。かつての議論のもやもやが思い出される。 https://www.python.org/dev/peps/pep-0342/ とはいえ、glossaryにもasyncioを匂わせた形で載っていて、その文意も、coroutineで実装されているみたいなニュアンスを汲み取るにとどめておく、みたいな形で見方が正しいと主張するのは無理筋そう。inspectモジュールにもcorountineへの言及が存在するし、asyncio内の文脈にとどめた用法と言うのは厳しそう。 https://docs.python.org/3/glossary.html ↩