with文の話3

with文専用のオブジェクトをより簡単に定義する contextlib.contextmanager を紹介します。

こちら で、print をハックする with文専用クラスを紹介しました。
このクラスを中身を見てみると、__enter__() と __exit__() という２つのメソッドが定義されているのみです。このようなクラスと同じ動作をする関数を作れるようにするのが contextlib.contextmanager です。 contextlib.contextmanager を使った実装例は以下のようになります。

import sys
import io
from contextlib import contextmanager

@contextmanager
def print_hack():
    sio = io.StringIO()
    tmp, sys.stdout = sys.stdout, sio
    yield sio
    sys.stdout = tmp

行なっていることは以前のクラスと同じです。変わった点は

• クラスから関数になった
• contextmanager のデコレータを使用する
• __enter__() と __exit__() に相当する部分が yield で区切られる。yield までが __enter__()、yield より後ろが __exit__()、となる

となります。

関数の使い方は以下の通りで、以前のクラスと全く同じです。

# with文でprint_hack()を呼び出す
>>> with print_hack() as hack:
...     print('Hello hack')
...
>>>   # ここではprintされない

# with内のprintの内容を確認
>>> hack.getvalue()
'Hello hack\n'

また、ここで定義した関数は通常の関数呼び出しからは呼び出すことはできません。誤使用を防ぐという意味で、大事な点です。

>>> print_hack()
<contextlib._GeneratorContextManager object at 0x01780470>

個人的には、クラスを定義すると self をたくさん書かなければならない点が苦痛です。関数で済むなら関数で済ます方が楽で良いのではないでしょうか。