Python: デコレータとコンテキストマネージャの両方として動作する API の作り方

以前、仕事でデータベースのトランザクションを管理する Python の API を考えているときに、同僚と「デコレータで作るべき」「いやコンテキストマネージャの方が扱いやすい」みたいなやり取りをしたことがあった。 でも、今考えるとどちらとしても動作するように作っておけばよかったんだよね。

ということで、今回はデコレータとしても動作するしコンテキストマネージャとしても動作するような API を作る方法について書くことにする。 補足しておくと、デコレータもコンテキストマネージャも対象となる処理の前後に特定の処理を挿入できるというところに共通点がある。 今回紹介する内容も、対象となる処理の前後に特定の処理を挿入するというものになっている。

サンプルコードその一

以下のサンプルコードでは、clamp() 関数がデコレータとしてもコンテキストマネージャとしても動作するように作られている。 clamp() 関数の内容は、対象となる処理の前後に print() 関数の実行を加えるというもの。 @clamp デコレータで修飾された decorator() 関数や、clamp() をコンテキストマネージャとして with 文で print() 関数をネストさせている contextmanager() 関数がそれに当たる。 とはいえ clamp() 関数自体は難しいことをやっているわけではなく、引数がある場合にはデコレータ (_decorator) を、ない場合にはコンテキストマネージャ (_contextmanager()) を返しているに過ぎない。 尚、コンテキストマネージャは本来であればクラスに特殊メソッド __enter__() と __exit__() を実装することになるが、今回は楽をするために @contextlib.contextmanager() でジェネレータを修飾することで作成している。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import contextlib


@contextlib.contextmanager
def _contextmanager():
    print('Beofre')
    try:
        yield
    finally:
        print('After')


def _decorator(func):
    def __decorator(*args, **kwargs):
        print('Beofre')
        try:
            ret = func(*args, **kwargs)
        finally:
            print('After')
            return ret
    return __decorator


def clamp(func=None):
    '''
    引数が指定されればデコレータとして、指定されなければコンテキストマネージャとして使える
    '''
    if func is not None:
        return _decorator(func)
    else:
        return _contextmanager()


# デコレータとして使うパターン
@clamp
def decorator():
    print('Hello, World!')


def contextmanager():
    # コンテキストマネージャとして使うパターン
    with clamp():
        print('Hello, World!')


def main():
    print('---デコレータ---')
    decorator()
    print('---コンテキストマネージャ---')
    contextmanager()


if __name__ == '__main__':
    main()

実行結果は次の通り。 思惑通り 'Hello, World!' の前後に 'Before' と 'After' が出力されている。

$ python clamp1.py
---デコレータ---
Beofre
Hello, World!
After
---コンテキストマネージャ---
Beofre
Hello, World!
After

サンプルコードその二

次はより実用的な例として、デコレータが修飾した関数の引数にオブジェクトをインジェクトしたり、コンテキストマネージャが as キーワードで受け取れるオブジェクトを返すようにしてみよう。 これは例えばデータベースのトランザクションを管理する API であれば、セッションを表すオブジェクトがそれに当たることになると思う。

次のサンプルコードでは、MyClass クラスのインスタンスをデコレータやコンテキストマネージャ経由でユーザに使えるようにしたい、というのを想定している。 コンテキストマネージャの場合はシンプルに yield で返したいオブジェクトを指定すれば、それが with 文の as キーワードに続く変数で受け取れるようになる。 デコレータの場合はそれよりも多少複雑で、まずデコレータの引数でインジェクト対象の変数名を受け取っている。 その上で、デコレータの中で修飾対象となる関数の引数を inspect モジュールを使って調べた上でオブジェクトを引数にインジェクトしている。 ちなみに以下のコードではキーワード付き変数にはインジェクトできない点には注意が必要。 デコレータの場合もコンテキストマネージャの場合も、受け取ったオブジェクトのメソッド (MyClass#echo()) を使ってメッセージを出力しているところがポイントになる。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import contextlib
import inspect


# デコレータで修飾された関数や with 文のスコープ内で使いたいオブジェクトを想定したクラス
class MyClass(object):

    def echo(self, msg):
        print(msg)


@contextlib.contextmanager
def _contextmanager():
    print('Beofre')
    try:
        # yield で返したオブジェクトが with 文で as に続く変数に代入される
        yield MyClass()
    finally:
        print('After')


def _decorator(name):
    def __decorator(func):
        def ___decorator(*args, **kwargs):
            print('Beofre')
            try:
                # デコレータでラップする対象の関数から引数名をリストで取得する
                args_names = inspect.getargspec(func)[0]
                # 名前を手がかりにオブジェクトをインジェクトすべき場所を見つける
                target_pos = args_names.index(name)
                # 元々の引数を配列にする (タプルは Immutable なオブジェクトのため)
                args_list = list(args)
                # 必要なオブジェクトをインジェクトする
                args_list.insert(target_pos, MyClass())
                # 加工後の引数を使って関数を呼び出す
                ret = func(*args_list, **kwargs)
            finally:
                print('After')
                return ret
        return ___decorator
    return __decorator


def clamp(name=None):
    '''
    引数が指定されればデコレータとして、指定されなければコンテキストマネージャとして使える
    '''
    if name is not None:
        return _decorator(name)
    else:
        return _contextmanager()


# デコレータとして使うパターン
@clamp('obj')  # オブジェクトをインジェクトする引数名を指定する
def decorator(obj):
    obj.echo('Hello, World!')


def contextmanager():
    # コンテキストマネージャとして使うパターン
    with clamp() as obj:
        obj.echo('Hello, World!')


def main():
    print('---デコレータ---')
    decorator()
    print('---コンテキストマネージャ---')
    contextmanager()


if __name__ == '__main__':
    main()

上記を実行してみよう。 実行結果自体は先程と変わらない。

$ python clamp2.py
---デコレータ---
Beofre
Hello, World!
After
---コンテキストマネージャ---
Beofre
Hello, World!
After

ばっちり。

やっぱり、かっこいい API を設計したいよね。

スマートPythonプログラミング: Pythonのより良い書き方を学ぶ