Python の contextlib.ExitStack で動的にリソースを管理する

contextlib.ExitStack は、動的に複数のコンテキストマネージャを管理するためのクラスです。リソースの数が事前にわからない場合や、条件によってリソースを追加する場合に便利です。

基本的な使い方

enter_context() でコンテキストマネージャを追加します。

from contextlib import ExitStack

with ExitStack() as stack:
    f1 = stack.enter_context(open("file1.txt", "w"))
    f2 = stack.enter_context(open("file2.txt", "w"))
    f1.write("Hello")
    f2.write("World")
# ブロック終了時に両方のファイルが閉じられる

追加した順番と逆の順序でリソースが解放されます。

動的なリソース管理

ファイルのリストを動的に処理する例です。

from contextlib import ExitStack

filenames = ["data1.txt", "data2.txt", "data3.txt"]

with ExitStack() as stack:
    files = [stack.enter_context(open(name, "w")) for name in filenames]
    
    for i, f in enumerate(files):
        f.write(f"Content {i}")
# すべてのファイルが自動で閉じられる

通常の with 文では、ファイル数が事前にわからない場合に対応できませんが、ExitStack なら可能です。

コールバックの登録

callback() で終了時に実行する関数を登録できます。

from contextlib import ExitStack

def cleanup(name):
    print(f"クリーンアップ: {name}")

with ExitStack() as stack:
    stack.callback(cleanup, "リソースA")
    stack.callback(cleanup, "リソースB")
    print("処理中")

# 出力:
# 処理中
# クリーンアップ: リソースB
# クリーンアップ: リソースA

コールバックは登録した順と逆の順序で実行されます。

push() でコンテキストマネージャを追加

push() を使うと、__exit__ だけを登録できます。

from contextlib import ExitStack

class Resource:
    def __init__(self, name):
        self.name = name
    
    def close(self):
        print(f"{self.name} を閉じました")

with ExitStack() as stack:
    r = Resource("DB接続")
    stack.callback(r.close)  # close() を終了時に呼ぶ
    print("処理中")
# 処理中
# DB接続 を閉じました

pop_all() でリソースを引き継ぐ

pop_all() を使うと、リソースの管理を別の ExitStack に移せます。

from contextlib import ExitStack

def setup_resources():
    with ExitStack() as stack:
        f1 = stack.enter_context(open("temp1.txt", "w"))
        f2 = stack.enter_context(open("temp2.txt", "w"))
        # 成功したらリソースを返す（閉じない）
        return stack.pop_all(), [f1, f2]

# リソースを受け取る
resource_stack, files = setup_resources()

# 後で明示的に閉じる
with resource_stack:
    for f in files:
        f.write("data")
# ここでファイルが閉じられる

条件付きリソース管理

条件によってリソースを追加するかどうかを決められます。

from contextlib import ExitStack

def process(use_cache=False, use_log=False):
    with ExitStack() as stack:
        if use_cache:
            cache = stack.enter_context(open("cache.txt", "w"))
        else:
            cache = None
        
        if use_log:
            log = stack.enter_context(open("log.txt", "w"))
        else:
            log = None
        
        # 処理
        if cache:
            cache.write("cached data")
        if log:
            log.write("log entry")

process(use_cache=True, use_log=True)

エラー時の部分クリーンアップ

複数リソースの取得中にエラーが発生した場合、取得済みのリソースだけをクリーンアップできます。

from contextlib import ExitStack

with ExitStack() as stack:
    try:
        f1 = stack.enter_context(open("file1.txt", "w"))
        f2 = stack.enter_context(open("file2.txt", "w"))
        # f2 の取得中にエラーが発生しても、f1 は閉じられる
    except IOError:
        print("ファイルオープンに失敗")

ExitStack は、リソース管理が複雑になる場合に非常に強力なツールです。