Python 標準ライブラリ contextlib with文用コンテキスト

Python標準のライブラリcontextlibを使うと、with文に対応したコンテキストを表現できます。

デコレータでコンテキストマネージャーを使用

ContextDecoratorを継承してクラスを宣言することで、コンテキストマネージャーとしてデコレータで使えるようになる。
非同期版のAsyncContextDecoratorもある。

from contextlib import ContextDecorator

class SampleContext(ContextDecorator):
    def __enter__(self):
        print('__enter__')
        return self

    def __exit__(self, *args):
        print('__exit__')
        return False

@SampleContext()
def sample_function():
    print('sample_function')

sample_function()

__enter__
sample_function
__exit__

with文から抜ける時にcloseしてくれるコンテキストマネージャー

closingで指定したクラスのcloseメソッドを終了時に呼び出してくれるコンテキストマネージャーを取得できる。
非同期版はaclosingを使用する。

from contextlib import closing

class SampleClass:
    def close(self):
        print('close')
    def hello(self):
        print('hello')

with closing(SampleClass()) as c:
    c.hello()

hello
close

関数をコンテキストマネージャーとして使えるようにする

@contextmanagerデコレーを付与することで、関数をコンテキストマネージャーとして使えるようにできる。
関数はyieldを返す必要がある。
非同期版は@asynccontextmanager。

from contextlib import contextmanager

@contextmanager
def sample_context():

    print('create sample context')

    try:
        yield 'yield here'
    finally:
        print('exit sample context')

with sample_context() as sc:
    print(sc)

create sample context
yield here
exit sample context

コンテキストマネージャーが取得できない場合

openに失敗した場合なども含め、with文で統一して処理を記述したいような時にnullcontextを使うことができる。

import contextlib

def get_context_or_nullcontext():

    f = None

    try:
        f = open('./sample.txt')
    except Exception as e:
        print(e)
    finally:
        if f is None:
            f = contextlib.nullcontext('enter_result')

    return f

with get_context_or_nullcontext() as c:
    print(c) # TextIOWrapper or 'enter_result'

標準出力先を変更

redirect_stdoutで標準出力先を変更できる。
以下はprintで標準出力ではなくsample.txtに内容が書き込まれる。

with open('./sample.txt', 'w') as f:
    with contextlib.redirect_stdout(f) as a:
        print('write!')

参考

• cpython/contextlib.py at 3.10 · python/cpython
• contextlib --- with 文コンテキスト用ユーティリティ — Python 3