Python の contextlib とジェネレータでコンテキストマネージャを作る
コンテキストマネージャを作るには通常 enter と exit を持つクラスを定義する。しかし contextlib の @contextmanager デコレータを使えば、ジェネレータ関数 1 つでコンテキストマネージャを作れる。yield の「一時停止と再開」という性質が、with 文の「前処理 → 本体 → 後処理」という構造とぴったり噛み合うためだ。
@contextmanager の基本
contextlib.contextmanager は、yield を 1 つだけ含むジェネレータ関数をコンテキストマネージャに変換するデコレータだ。
from contextlib import contextmanager
@contextmanager
def tag(name):
print(f"<{name}>")
yield
print(f"</{name}>")
with tag("div"):
print("コンテンツ")実行結果は以下のようになる。
# <div>
# コンテンツ
# </div>yield より前が enter に相当し、yield より後が exit に相当する。with ブロックの本体は yield で一時停止している間に実行される。
yield で値を渡す
yield に値を添えると、with 文の as 節でその値を受け取れる。
from contextlib import contextmanager
@contextmanager
def open_db(connection_string):
print(f"接続: {connection_string}")
conn = {"status": "connected", "queries": 0}
yield conn
print(f"切断: {conn['queries']} 件のクエリを実行")
with open_db("localhost:5432/mydb") as db:
db["queries"] += 1
print(f"状態: {db['status']}")
db["queries"] += 1yield conn の conn が as db に渡される。yield の後に書いた切断処理は、with ブロックを抜けたときに実行される。
クラスベースとの比較
同じコンテキストマネージャをクラスで書いた場合と並べてみる。
enter と exit を定義する正統な方法。状態管理が複雑な場合やメソッドを追加したい場合に向く。ただしボイラープレートが多く、前処理と後処理がメソッドに分断されるため流れを追いにくい。
ジェネレータ関数 1 つで完結する。yield の前後が前処理・後処理になるため、処理の流れが上から下へ一直線に読める。シンプルなリソース管理に最適。
# クラスベース
class Timer:
def __enter__(self):
import time
self.start = time.time()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
import time
self.elapsed = time.time() - self.start
print(f"{self.elapsed:.3f} 秒")
return False
# @contextmanager ベース
from contextlib import contextmanager
import time
@contextmanager
def timer():
start = time.time()
yield
elapsed = time.time() - start
print(f"{elapsed:.3f} 秒")クラスベースでは self を介して start を受け渡す必要があるが、ジェネレータ版ではローカル変数がそのまま yield をまたいで生き続ける。これがジェネレータの一時停止・再開という仕組みの恩恵だ。
例外処理を組み込む
with ブロック内で例外が発生すると、yield の位置で例外が再送出される。try/finally で囲めば、例外の有無にかかわらず後処理を確実に実行できる。
from contextlib import contextmanager
@contextmanager
def managed_resource(name):
print(f"リソース {name} を確保")
try:
yield name
finally:
print(f"リソース {name} を解放")
try:
with managed_resource("GPU") as res:
print(f"{res} を使用中")
raise RuntimeError("メモリ不足")
except RuntimeError as e:
print(f"エラー捕捉: {e}")with ブロック内で例外が発生
yield の位置で例外が再送出される
try/finally の finally 節が実行される
例外は with ブロックの外へ伝播する
finally を使わず except で例外を握りつぶすこともできるが、一般的には finally で後処理を行い、例外はそのまま伝播させるのが安全だ。
例外を抑制するパターン
特定の例外を意図的に無視したい場合は、except で捕捉し yield の後に処理を続ければよい。
from contextlib import contextmanager
@contextmanager
def ignore_error(*exceptions):
try:
yield
except exceptions as e:
print(f"無視: {type(e).__name__}: {e}")
with ignore_error(FileNotFoundError, PermissionError):
open("存在しないファイル.txt")
print("処理を続行")except で例外を捕捉してジェネレータが正常に終了すると、with ブロックの外には例外が伝播しない。これは contextlib.suppress と同じ動作だ。
ジェネレータ内で例外を捕捉して yield せずに終了すると、exit が True を返したのと同じ扱いになり、例外が抑制される。
exit(exc_type, exc_val, exc_tb) が True を返すと例外が伝播しなくなる仕組みと同等。
yield は 1 つだけ
@contextmanager で装飾するジェネレータ関数には yield を 1 つだけ書く。複数の yield があると RuntimeError になる。
from contextlib import contextmanager
@contextmanager
def bad_context():
yield "first"
yield "second" # これが問題
try:
with bad_context() as val:
print(val)
except RuntimeError as e:
print(f"エラー: {e}")
# エラー: generator didn't stopwith 文は「入って出る」という 1 回の遷移しか持たないため、yield が 2 つ以上あると「終了すべきタイミングで終了しなかった」として RuntimeError が発生する。
実用例: 一時的な設定変更
グローバルな設定を一時的に変更し、with ブロックを抜けたら元に戻すパターンはよく使われる。
from contextlib import contextmanager
import os
@contextmanager
def env_var(key, value):
old = os.environ.get(key)
os.environ[key] = value
try:
yield value
finally:
if old is None:
del os.environ[key]
else:
os.environ[key] = old
with env_var("DEBUG", "true") as v:
print(f"DEBUG = {os.environ['DEBUG']}") # true
print(f"DEBUG = {os.environ.get('DEBUG', '未設定')}") # 元に戻る環境変数の書き換え、ロケールの切り替え、ログレベルの変更など、「一時的に状態を変えて確実に戻す」操作にこのパターンは適している。
実用例: ディレクトリの一時移動
from contextlib import contextmanager
import os
@contextmanager
def cd(path):
prev = os.getcwd()
os.chdir(path)
try:
yield path
finally:
os.chdir(prev)
with cd("/tmp"):
print(os.getcwd()) # /tmp
print(os.getcwd()) # 元のディレクトリに戻る前処理でディレクトリを移動し、後処理で元に戻す。try/finally で囲んでいるため、with ブロック内で例外が発生してもカレントディレクトリは確実に復元される。
ジェネレータの仕組みとの関係
@contextmanager の内部では、ジェネレータプロトコルがそのまま活用されている。
デコレータが内部でジェネレータの next() を呼ぶ。ジェネレータは yield まで進み、yield の値が as 節に渡される。
ジェネレータは yield で一時停止したまま待機する。この間に with ブロック内のコードが実行される。
正常終了なら再度 next() が呼ばれ、ジェネレータは yield の後から再開する。例外発生時は throw() で例外がジェネレータに送り込まれる。
send() や throw() といったジェネレータの制御メソッドが、コンテキストマネージャの入口と出口にきれいに対応している。ジェネレータの一時停止と再開という仕組みが、リソース管理の「確保→使用→解放」という流れを自然に表現できる理由はここにある。
@contextmanager を使ったジェネレータ関数で、with ブロック内で例外が発生した場合、ジェネレータ内部では何が起きるか？
- ジェネレータが自動的に終了し、finally は実行されない
- yield の位置で例外が再送出され、try/except や finally で処理できる
- 例外は自動的に握りつぶされ、yield の後が通常通り実行される
- StopIteration に変換されてジェネレータが終了する
@contextmanager で装飾したジェネレータ関数に yield を 2 つ書くとどうなるか？
- 2 回目の yield の値は無視される
- with 文が 2 回実行される
- RuntimeError が発生する
- StopIteration が発生する
with 文は 1 回の「入って出る」遷移しか持たない。exit で next() を呼んだとき、ジェネレータが終了せず 2 つ目の yield に到達すると generator didn’t stop という RuntimeError が送出される。
with ブロック内の例外は throw() によって yield の位置に送り込まれる。ジェネレータ内で try/finally を書いておけば、例外の有無にかかわらず後処理を確実に実行できる。