withで使える関数を作る方法 | withについて学ぶ [python]

withとは

最初にwithの具体的な例を挙げる
たとえば、ファイルを開けたら責任をもってcloseしないといけない。

withを使わない場合の例

f = open("example.txt", "r")

print(f.read())

# 自分で閉じないといけない
f.close()

withを使う場合の例

# withを使うとブロックを抜けたときに閉じてくれる
with open("example.txt", "r") as f:
    print(f.read())

では、このようにブロックに入るとき、出るときに特定の処理をさせる方法を解説します。

ちなみにopen関数は組み込み関数で、C言語で書かれているため以下で紹介するのとは厳密には違うと思いますが、例として出すにはわかりやすそうだったので出しました。

実際に作ってみる

公式ドキュメントを参考にしています。

クラスの場合

関数と比べて可読性が良く、分かりやすいです。

class Context:
    def __enter__(self):
        print("enter")

    def __exit__(self, *args):
        print("exit")


if __name__ == "__main__":
    with Context():
        print("body")

実行結果

enter
body
exit

解説

ね、意外と簡単でしょ

__enter__メソッドはブロックに入る前に実行される処理を定義します。
__exit__はwithブロックを抜けたときの処理を定義します。

with文の進行について　

引用：以下リンク

一つの "要素" を持つ with 文の実行は以下のように進行します:

コンテキスト式 (with_item で与えられた式) を評価することで、コンテキストマネージャを取得します。

コンテキストマネージャの enter() メソッドが、後で使うためにロードされます。

コンテキストマネージャの exit() メソッドが、後で使うためにロードされます。

コンテキストマネージャの enter() メソッドが呼ばれます。

with 文にターゲットが含まれていたら、それに enter() からの戻り値が代入されます。

asを使うには

class FileContext:
    def __init__(self, file_name: str, mode):
        self.f = open(file_name, mode)

    def __enter__(self):
        return self.f

    def __exit__(self, *args):
        self.f.close()


if __name__ == "__main__":
    with FileContext("example.txt", "r") as f:
        print(f.read())

インスタンスを返したい場合は__enter__のreturnに書けばいいだけです。
これにより、as で任意の名前をつけてインスタンスがブロック内で使用できます。

(open関数では、すでにこのような実装がされているのでこんな事をしなくても使えます。あくまでも例です。)

実用例として考えられるのはクラスを作ってそこで毎回呼ぶcloseのような処理が必要だとか、ブロックのはじめから終わりまでの線をprintしたいときに出番が来ると思います。

関数の場合

クラスを使う場合と比べて、少し面倒だと感じます。

最低限必要なコード

最低限tryとfinallyそしてyieldは実装する必要があります。

from contextlib import contextmanager


# デコレーターを使う
@contextmanager
def WithContext():
    try:
        yield print("try")
    finally:
        print("finally")


if __name__ == "__main__":
    with WithContext():
        print("Context")

実行結果

try
Context
finally

詳しい解説

ここからは内部的にどうなっているかなどの結構深い話をします

from contextlib import contextmanager


# デコレーターを使う
@contextmanager
def FileContext(file: str, mode="r"):
    # このブロックに入ったときに実行したい処理を書く
    f = open(file, mode)

    # このyieldがreturnのようなもの
    try:
        yield f

    # ResutlContext()の中で発生した例外はここで処理する
    except Exception as e:
        pass

    # ResutlContext()の中で発生した例外があってもなくても必ず実行する
    finally:
        f.close()


if __name__ == "__main__":
    with FileContext("example.txt") as file:
        print(file.read())

まず、@contextmanagerというデコレータを使います。
これにより、withブロックを抜けたときにfinallyで定義した処理が実行されるようになります。
yieldがreturnのように働きます。
yieldは「イールド」と読みます。

なぜ、yieldが必要なのか

デコレートをつけた関数は呼び出されたときに デコレータ下の関数を取り込んでデコレータ内で実行することができます。(実際にデコレータを作ってみると意味が分かると思います。)

デコレータ内では_GeneratorContextManagerクラスを呼び出します。

_GeneratorContextManagerメソッドの__enter__関数内で実行されるのはnext(self.gen)関数です、このself.genは継承元である_GeneratorContextManagerBaseで受け取った関数をself.genに代入しています、つまりここに@デコレータをつけた関数が入ります。ここで言うとFileContext関数です。

next()関数は iterator型の__next__()メソッドを呼び出す関数です。
(iteratorとは繰り返し可能なものって認識でいいと思います。)

なので、next()に渡す関数はiterator型である必要があります。(繰り返し可能であること)
そこで、yieldの出番です、yieldは関数の本体で yield式 を使用するとその関数はジェネレータ(generator)を返すようになります。
ジェネレータはイテレータ(iterator)を返します。
（厳密にはgenerator iteratorを返すらしい）

なので、yieldが必要なのです。

_GeneratorContextManagerBase の一部

class _GeneratorContextManagerBase:
    """Shared functionality for @contextmanager and @asynccontextmanager."""

    def __init__(self, func, args, kwds):
        self.gen = func(*args, **kwds)
        self.func, self.args, self.kwds = func, args, kwds

_GeneratorContextManager の一部

class _GeneratorContextManager(
    _GeneratorContextManagerBase,
    AbstractContextManager,
    ContextDecorator,
):
    """Helper for @contextmanager decorator."""

    def __enter__(self):
        # do not keep args and kwds alive unnecessarily
        # they are only needed for recreation, which is not possible anymore
        del self.args, self.kwds, self.func
        try:
            return next(self.gen) # ここでジェネレータイテレータでないといけない
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

解説は以上になります。
デコレータを使う方法よりクラスでenter,exitのメソッドを使えるようになるとよりスマートなプログラムが書けるのではないかと思います。
ぜひ、覚えてみてください。