Pythonで独自のコンテキストマネージャを作成する

Pythonでは、withブロックで使用できるオブジェクトをコンテキストマネージャと呼びます。コンテキストマネージャは、__enter__メソッドと__exit__メソッドを持つオブジェクトを作成することで作成できます。

class Example:
    def __enter__(self):
        print("enter")
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("exit")

このコンテキストマネージャは、withブロックが開始されるときに”enter”を出力し、withブロックが終了するときに”exit”を出力します。

with Example():
    print("Yay Python!")

出力:

enter
Yay Python!
exit

次に、実際に何か役立つことをするコンテキストマネージャを見てみましょう。このコンテキストマネージャは、一時的に環境変数の値を変更します。

import os

class set_env_var:
    def __init__(self, var_name, new_value):
        self.var_name = var_name
        self.new_value = new_value

    def __enter__(self):
        self.original_value = os.environ.get(self.var_name)
        os.environ[self.var_name] = self.new_value

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.original_value is None:
            del os.environ[self.var_name]
        else:
            os.environ[self.var_name] = self.original_value

このコンテキストマネージャを使用すると、そのwithブロック内で、USER環境変数は異なる値を持ちます。

with set_env_var("USER", "akin"):
    print("USER env var is", os.environ["USER"])

出力:

USER env var is akin

しかし、withブロックが終了した後、その環境変数の値は元の値に戻ります。

print("USER env var is", os.environ["USER"])

出力:

USER env var is trey

これはすべて、コンテキストマネージャの__enter__メソッドと__exit__メソッドのおかげで、コンテキストマネージャのwithブロックが開始されたときと終了したときに実行されます。