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()が存在する1。コレを使えば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) -> None:
        ...

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

gist


  1. throwの存在は、methaneさんに教えてもらった。 https://mobile.twitter.com/methane/status/1325302752944160768

  2. ちなみに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

pythonでbreakpointなしに例外の発生元で自動的にpdbする方法のメモ

デバッガーを使いたいがコードにbreakpoint を仕込むのが面倒な時がある。

例えば以下のようなコードがあるとして、 foo > bar > boo と辿っていた中でのbooでpdbを実行して欲しい。

def foo():
    print("foo")
    bar()
    print("foo")


def bar():
    print("bar")
    boo()
    print("bar")


def boo():
    print("boo")
    raise Exception("oops") # ここでpdbを

foo()

oopsの部分に、breakpointなしで、pdbしたい。

回答

正解から言うと、以下のような関数を作れば良いらしい1

すごく古いオライリーのクックブックに載っていた(https://www.oreilly.com/library/view/python-cookbook/0596001673/ch14s06.html)。

def info(type, value, tb):
    if hasattr(sys, "ps1") or not sys.stderr.isatty():
        # You are in interactive mode or don't have a tty-like
        # device, so call the default hook
        sys.__excepthook__(type, value, tb)
    else:
        import traceback, pdb

        # You are NOT in interactive mode; print the exception...
        traceback.print_exception(type, value, tb)
        # ...then start the debugger in post-mortem mode
        pdb.pm()


sys.excepthook = info

intractive shellなどへの対応を無視しちゃうと実質これだけ。

import traceback, pdb


def info(type, value, tb):
    pdb.pm()


sys.excepthook = info

詳細

以下詳細。キモは sys.excepthook による補足と、 pdbのpost moterm mode。

sys.excepthook

sys.excepthookが使える。

例外が発生し、その例外が捕捉されない場合、インタプリタは例外クラス・例外インスタンス・トレースバックオブジェクトを引数として sys.excepthook を呼び出します。対話セッション中に発生した場合はプロンプトに戻る直前に呼び出され、Pythonプログラムの実行中に発生した場合はプログラムの終了直前に呼び出されます。このトップレベルでの例外情報出力処理をカスタマイズする場合、 sys.excepthook に引数の数が三つの関数を指定します。

まぁ要は以下のようなsignatureの関数を受け取って、補足されない例外に対してこのhookが効くということのようだ2

def hook(typ, value, tb):
    ...

このhookに pdbをうまくしかけられると良い。

pdb.postmortem

実は、pdb.post_mortem()というメソッドがあることを知らなかったのだけれど、コレはその名の通り事後の解析をするためのメソッドのようだ。

内部的にはtracebackオブジェクトを渡せばそれが使われるし、そうでない場合は、sys.exc_infoから現在補足している例外を取り出して利用する。実際実装は以下のようになっていた3

def post_mortem(t=None):
    # handling the default
    if t is None:
        # sys.exc_info() returns (type, value, traceback) if an exception is
        # being handled, otherwise it returns None
        t = sys.exc_info()[2]
    if t is None:
        raise ValueError("A valid traceback must be passed if no "
                         "exception is being handled")

    p = Pdb()
    p.reset()
    p.interaction(None, t)

def pm():
    post_mortem(sys.last_traceback)

(実は、reset()や諸々でどのようにpdbが動いているかを知ろうとすると結構面白かったりする。内部的には bdbモジュールが使われている4)

実際しっかり止まる

$ python 03*.py
foo
bar
boo
Traceback (most recent call last):
  File "03pdb-on-exception.py", line 37, in <module>
    foo()
  File "03pdb-on-exception.py", line 6, in foo
    bar()
  File "03pdb-on-exception.py", line 12, in bar
    boo()
  File "03pdb-on-exception.py", line 18, in boo
    raise Exception("oops")
Exception: oops
> 03pdb-on-exception.py(18)boo()
-> raise Exception("oops")
(Pdb) q

素晴らしいですね。

gist


  1. shimizukawaさんに教えてもらった https://mobile.twitter.com/shimizukawa/status/1326498954574356481

  2. この3つの引数は context managerを作るときの __exit__() などでもおなじみ

  3. 実装を考えると、クックブックのコードよりは、tracebackを直接渡すようなコードにしたほうがきれいな気もする。

  4. コレを悪用すると、withの中を評価せずに、ソースコードを取り出し、変換を加えてexecするみたいな事もできたりする。 https://gist.github.com/podhmo/b2aaf26e14999c871509783fe8f55909