Python の contextlib.contextmanager でコンテキストマネージャを作る

contextlib.contextmanager デコレータを使うと、クラスを定義せずにジェネレータ関数からコンテキストマネージャを作成できます。シンプルなコンテキストマネージャを手軽に作りたい場合に便利です。

基本的な使い方

@contextmanager デコレータを付けたジェネレータ関数を定義します。

from contextlib import contextmanager

@contextmanager
def my_context():
    print("開始")
    yield
    print("終了")

with my_context():
    print("処理中")

# 出力:
# 開始
# 処理中
# 終了

yield の前が __enter__() に、後が __exit__() に相当します。

yield で値を返す

yield で値を返すと、as で受け取れます。

from contextlib import contextmanager

@contextmanager
def open_file(filename, mode):
    f = open(filename, mode)
    try:
        yield f
    finally:
        f.close()

with open_file("data.txt", "w") as f:
    f.write("Hello")

try-finally パターン

例外が発生してもクリーンアップを確実に行うには、try-finally を使います。

from contextlib import contextmanager

@contextmanager
def managed_resource():
    print("リソース確保")
    resource = {"data": []}
    try:
        yield resource
    finally:
        print("リソース解放")
        resource["data"].clear()

with managed_resource() as r:
    r["data"].append(1)
    raise ValueError("エラー")  # finally は実行される

クラスベースとの比較

クラスベース

状態管理が複雑な場合に適している。__enter__ と __exit__ を明示的に定義。コードが長くなりがち。

@contextmanager

シンプルなケースに適している。ジェネレータで直感的に書ける。コードが簡潔。

同じ機能をクラスで書くと以下のようになります。

# クラスベース
class Timer:
    def __enter__(self):
        self.start = time.time()
        return self
    
    def __exit__(self, *args):
        print(f"経過: {time.time() - self.start:.2f}秒")

# @contextmanager
@contextmanager
def timer():
    start = time.time()
    yield
    print(f"経過: {time.time() - start:.2f}秒")

例外の処理

try-except で例外をキャッチして処理することもできます。

from contextlib import contextmanager

@contextmanager
def suppress_value_error():
    try:
        yield
    except ValueError as e:
        print(f"ValueError を抑制: {e}")

with suppress_value_error():
    raise ValueError("テスト")
print("処理継続")
# ValueError を抑制: テスト
# 処理継続

例外をキャッチしなければ、そのまま外部に伝播します。

引数を受け取る

通常の関数と同様に引数を受け取れます。

from contextlib import contextmanager

@contextmanager
def indent(level):
    prefix = "  " * level
    print(f"{prefix}開始")
    yield prefix
    print(f"{prefix}終了")

with indent(2) as prefix:
    print(f"{prefix}処理中")

# 出力:
#     開始
#     処理中
#     終了

@contextmanager を使えば、数行でコンテキストマネージャを作成できます。シンプルなリソース管理や状態の一時変更に最適です。