Python: contextlib.contextmanager を使うときの注意点

Python には with 文というシンタックスがあって、これを使うと前処理と後処理が書きやすくなっている。その際、with 文にはコンテキストマネージャという種類のオブジェクトを渡すことが取り決めになっている。

そのコンテキストマネージャは標準ライブラリの contextlib.contextmanager を使うと楽に作ることができる。しかし、それを使うにはちょっとした注意点があるというのが今回の本題になる。

コンテキストマネージャのおさらい

まずはコンテキストマネージャとは何か、というおさらいから。コンテキストマネージャは、__enter__() と __exit__() というふたつの特殊メソッドを実装したクラス (またはそのインスタンス) を指している。

以下は最も単純と思われるコンテキストマネージャを実装したサンプルコード。 MyManager クラスがコンテキストマネージャになっている。

#!/usr/bin/env python
# -*- coding: utf-8 -*-


class MyManager(object):

    def __enter__(self):
        # 最初に実行される (前処理)
        print('enter')

    def __exit__(self, exc_type, exc_value, traceback):
        # 最後に実行される (後処理)
        print('exit')


def main():
    with MyManager():
        # 前処理と後処理の間に実行される
        print('Hello, World!')


if __name__ == '__main__':
    main()

上記を実行すると、with 文の中の処理が __enter__() と __exit__() の処理に囲われて実行されることがわかる。

$ python contextmanager1.py
enter
Hello, World!
exit

contextlib.contextmanager を使ってみる

それでは、今度は contextlib.contextmanager を使った形でコンテキストマネージャを実装してみよう。ジェネレータを @contextlib.contextmanager デコレータで修飾すると、それがコンテキストマネージャになる。ジェネレータというのは return の代わりに yield で値を返すようにした関数のこと。 yield の箇所が with 文のブロック内の処理に置き換わって、前後が __enter__() と __exit__() に相当するようなイメージをもつと良い。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import contextlib


@contextlib.contextmanager
def mymanager():
    print('enter')
    yield
    print('exit')


def main():
    with mymanager():
        print('Hello, World!')


if __name__ == '__main__':
    main()

上記を実行してみる。先ほどと同じ結果になった。

$ python contextmanager2.py
enter
Hello, World!
exit

落とし穴

めでたしめでたし…と言いたいところだけど、実は上記には大きな問題がある。以下のように、with 文のブロック内で例外が上がった場合にどういった挙動になるだろうか？

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import contextlib


@contextlib.contextmanager
def mymanager():
    print('enter')
    yield
    print('exit')


def main():
    with mymanager():
        raise Exception('Oops!')


if __name__ == '__main__':
    main()

実行してみると yield の前の処理は実行されているものの、その後の処理が実行されていないことがわかる。コンテキストマネージャには、例外が上がった場合にも後処理が必ず実行されるようにする使われ方が多いのでこれはまずい。

$ python contextmanager3.py
enter
Traceback (most recent call last):
  File "contextmanager3.py", line 20, in <module>
    main()
  File "contextmanager3.py", line 16, in main
    raise Exception('Oops!')
Exception: Oops!

問題に対処する

上記の問題を解決するには yield の箇所で上がる例外を try で補足するようにした上で、後処理を必ず実行されるように finally のブロック内に書く必要がある。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import contextlib


@contextlib.contextmanager
def mymanager():
    print('enter')
    try:
        yield
    finally:
        print('exit')


def main():
    with mymanager():
        raise Exception('Oops!')


if __name__ == '__main__':
    main()

上記であれば、例外になったとしても後処理が必ず実行される。

$ python contextmanager4.py
enter
exit
Traceback (most recent call last):
  File "contextmanager4.py", line 22, in <module>
    main()
  File "contextmanager4.py", line 18, in main
    raise Exception('Oops!')
Exception: Oops!

めでたしめでたし。