デコレータ・ジェネレータ・with ― 関数と処理を拡張する

導入

ここでは、実務のコードでよく見かける3つの「ちょっと高度な」機能をまとめて見ていきます。関数に処理を追加するデコレータ、値を1つずつ返すジェネレータ、そして第7章で使い方だけ紹介した with を自分で作る方法です。

説明

デコレータは、既存の関数を変更せずに、その前後に処理を付け足す仕組みです。@デコレータ名 を関数の直前に書きます。

flowchart LR
  call["呼び出し側"] --> wrap["wrapper(デコレータが用意した関数)"]
  wrap -->|"前処理"| pre["print など"]
  pre --> orig["元の関数本体"]
  orig --> post["後処理"]
  post --> call

def log_calls(fn):
    def wrapper(*args, **kwargs):
        print(f"呼び出し: {fn.__name__}{args}")
        return fn(*args, **kwargs)
    return wrapper

@log_calls
def add(a, b):
    return a + b

print(add(2, 3))

@log_calls を付けた瞬間、add は「log_calls が返した wrapper」に置き換わります。add(2, 3) を呼ぶと、実際には wrapper(2, 3) が呼ばれ、その中で元の add を呼び出す前後に好きな処理（ここではログ出力）を挟めます。

ジェネレータは、return の代わりに yield を使う関数です。呼び出しても即座に全部計算せず、値を求められるたびに1つずつ返します。

def countdown(n):
    while n > 0:
        yield n
        n -= 1

for v in countdown(3):
    print(v)

print(list(countdown(5)))

yield に達すると、そこで処理が一時停止して値を返します。次に値が求められると、止まった場所から再開します。大きな範囲を扱うときも、全部を一度にメモリに持たずに済むのが利点です。

最後に、with を自分のクラスで使えるようにする方法です。クラスに __enter__ と __exit__ という特殊メソッドを定義すると、そのクラスはwith に対応したコンテキストマネージャになります。

class FileLikeLogger:
    def __enter__(self):
        print("--- ログ開始 ---")
        return self
    def write(self, msg):
        print("LOG:", msg)
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("--- ログ終了 ---")
        return False

with FileLikeLogger() as logger:
    logger.write("処理を実行しました")

with FileLikeLogger() as logger: に入るとき __enter__ が、ブロックを抜けるとき（正常終了でもエラーでも）__exit__ が自動的に呼ばれます。第7章で見た「ファイルを with で開くと自動的に閉じる」仕組みは、まさにこの __enter__ / __exit__ によって実現されています。

やってみよう

log_calls を別の関数に付けてみたり、countdown の範囲を変えたり、FileLikeLogger の write を何度も呼んでみて、動きを確かめましょう。

演習

yield を使ったジェネレータ関数 countdown(n) を定義してください。n から 1 まで1つずつ減らしながら値を返すようにし、list(countdown(5)) の結果を表示してください（結果は [5, 4, 3, 2, 1]）。