Python で enter と exit を実装する

自分でコンテキストマネージャを作るには、クラスに __enter__() と __exit__() メソッドを実装します。これにより、with 文で使えるオブジェクトを定義できます。

基本的な構造

コンテキストマネージャには2つの特殊メソッドが必要です。

__enter__(self)

with ブロックに入る際に呼ばれる。as で受け取る値を返す。

__exit__(self, exc_type, exc_val, exc_tb)

with ブロックを抜ける際に呼ばれる。例外情報を受け取る。

シンプルな例

処理の開始と終了をログ出力するコンテキストマネージャを作ってみましょう。

class Logger:
    def __init__(self, name):
        self.name = name
    
    def __enter__(self):
        print(f"[{self.name}] 開始")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"[{self.name}] 終了")

with Logger("処理A"):
    print("何かの処理")

# 出力:
# [処理A] 開始
# 何かの処理
# [処理A] 終了

enter の戻り値

__enter__() の戻り値は as の変数に代入されます。

class Connection:
    def __enter__(self):
        print("接続を確立")
        return self  # self を返すのが一般的
    
    def __exit__(self, *args):
        print("接続を切断")
    
    def query(self, sql):
        print(f"実行: {sql}")

with Connection() as conn:
    conn.query("SELECT * FROM users")

# 出力:
# 接続を確立
# 実行: SELECT * FROM users
# 接続を切断

戻り値は self でなくても構いません。

class FileOpener:
    def __init__(self, filename):
        self.filename = filename
        self.file = None
    
    def __enter__(self):
        self.file = open(self.filename, "r")
        return self.file  # ファイルオブジェクトを返す
    
    def __exit__(self, *args):
        self.file.close()

with FileOpener("data.txt") as f:
    print(f.read())

exit の引数

__exit__() は3つの引数で例外情報を受け取ります。

class DebugContext:
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            print(f"例外の型: {exc_type}")
            print(f"例外の値: {exc_val}")
            print(f"トレースバック: {exc_tb}")
        return False  # 例外を再送出

with DebugContext():
    raise ValueError("テストエラー")

例外が発生しなかった場合、3つの引数はすべて None になります。

リソース管理の例

ファイルロックを管理するコンテキストマネージャの例です。

import os

class FileLock:
    def __init__(self, filename):
        self.lock_file = filename + ".lock"
    
    def __enter__(self):
        if os.path.exists(self.lock_file):
            raise RuntimeError("ファイルはロック中です")
        with open(self.lock_file, "w") as f:
            f.write("locked")
        return self
    
    def __exit__(self, *args):
        if os.path.exists(self.lock_file):
            os.remove(self.lock_file)

with FileLock("data.txt"):
    print("ロックを取得して処理中")
# ブロック終了でロック解放

クラスベースのコンテキストマネージャは、状態を持つ複雑なリソース管理に適しています。