Python Tips: 関数全体をコンテキストマネージャでかんたんにラップしたい

Python には with 構文で使える「コンテキストマネージャ」という種類のオブジェクトがあります。

with open('log.txt', 'w') as f:
    f.write('これはログです。')

この open() の戻り値の f がコンテキストマネージャです。

このコンテキストマネージャを使って関数をまるごとラップする方法を今回はご紹介してみたいと思います。

結論は「デコレータを使う」です。以下に具体例をあげてご説明していきます。

例として、 Fabric3 ( Python 3 に対応した Fabric のフォーク版)を使った、次のようなケースを考えてみましょう。

fabfile.py:

from fabric.api import cd, run, task

PROJECT_DIR = '/path/to/the/project/root'

@task
def git_status():
    '''`git status` を実行する'''
    with cd(PROJECT_DIR):
        run('git status')

@task
def git_pull():
    '''`git pull` を `master` ブランチに対して実行する'''
    with cd(PROJECT_DIR):
        run('git pull origin master')

ここでは git_status()git_pull() という 2 つの Fabric タスクが定義されています。いずれのタスクも、 PROJECT_DIR で表されたプロジェクトのルートディレクトリで処理を実行する前提となっています。そのために、 fabric.api.cd() を使って with cd(PROJECT_DIR) でコンテキストを作ってその中で実際の処理を実行する形になっています。

この with cd(PROJECT_DIR) の部分は共通です。タスクが 2 つの場合はそれほど気になりませんが、同じ with を使うタスクが増えてくるとこれを共通化したくなることでしょう。後々タスクが増えてきたときのためにコンテキストマネージャを共通化することを考えてみましょう。

結論としては上述のとおり、「関数全体をコンテキストマネージャでラップするデコレータを作る」形がよいかと思います。

実際にデコレータを導入してみた場合のコードは次のとおりとなります。

from functools import wraps
from fabric.api import cd, run, task

PROJECT_DIR = '/path/to/the/project/root'

def decorate_cd(directory):
    '''`fabric.api.cd()` で関数をラップするデコレータを生成する'''
    def decorator(func):
        '''関数をラップするデコレータ'''
        @wraps(func)
        def wrapper(*args, **kwargs):
            with cd(directory):
                func(*args, **kwargs)

        return wrapper

    return decorator

@task
@decorate_cd(PROJECT_DIR)
def git_status():
    '''`git status` を実行する'''
    run('git status')

@task
@decorate_cd(PROJECT_DIR)
def git_pull():
    '''`git pull` を `master` ブランチに対して実行する'''
    run('git pull origin master')

ポイントはもちろん、デコレータを生成する decorate_cd() 関数です。関数が 3 重に入れ子になっていてなんだかややこしそうに見えるので、外側の関数から順にご説明します。

decorate_cd() これは、ディレクトリのパスを受け取ってデコレータとなる関数を返す関数です。この関数自体はデコレータではなく、あくまでも「デコレータ関数を生成する関数」です。

decorator()decorate_cd() によって生成され、実際に関数に対してデコレート処理を行う関数です。関数に対するデコレータは「ただひとつの引数として関数を受け取って戻り値として関数を返す」ものである必要がありますが、この関数は func を引数として受け取り、 wrapper を戻りとして返しています。

wrapper() これは、実際に fabric.api.cd() のコンテキストマネージャで関数をラップする実際にラッパー関数です。

この構造で作られた decorate_cd() を使うことで、関数をかんたんにコンテキストマネージャでラップすることができるようになります。

デコレータで装飾された関数

@decorate_cd(PROJECT_DIR)
def git_status():
    '''`git status` を実行する'''
    run('git status')

は Python では

def git_status():
    '''`git status` を実行する'''
    run('git status')

git_status = decorate_cd(PROJECT_DIR)(git_status)

と同じ意味になるので、 git_status() 関数が無事に fabric.api.cd() でラップできていることになります。

うまく使えると便利かと思います。以上、関数全体をコンテキストマネージャでかんたんにラップする方法についてでした。

参考