Pythonのデコレータで関数に処理を追加する

はじめに

#

Pythonには、ある関数の前後に処理を追加する仕組みとして、デコレータと呼ばれる機能がある。本記事ではデコレータの使用例として、回帰の精度評価でMSE, RMSEの両方を算出する場合に、入力配列の型を変換する処理をデコレータにまとめたものを示す。

なお、本記事の内容とは直接関係ないが、WebアプリケーションライブラリのFlaskでコールバック関数にデコレータが使われていたことが勉強の動機である。

環境は以下の通り。

デコレータの例

#

デコレータの簡単な例を以下に示す。

def decorator(func):
    def wrapper(*args, **kwargs):
        print('--start--')
        func(*args, **kwargs)
        print('--end--')
    return wrapper

@decorator
def hello():
    print('Hello World')

hello()

実行結果

--start--
Hello World
--end--

上記の例では、decoratorがデコレータであり、helloが処理を追加される関数である。hello関数を定義する直前に@（デコレータ名）とすると、helloを呼び出したときにデコレータも同時に呼び出される。

デコレータ関数decoratorの内部にはさらに関数wrapperが定義されており、decoratorはwrapper関数のオブジェクトを返す。

wrapper関数では、func(*args, **kwargs)の部分でデコレートされる関数を実行している。その前後に実行したい処理を記述する。

デコレータの機能

#

デコレータには以下の使い方もある。

• デコレータを多重にネストできる
• デコレータが引数をとる

デコレータの使用例

#

デコレータには、以下のような使い道がある。

• 関数の実行時間を計測する
• 関数の入力値をチェックする
• 関数をいつ呼び出したか等の実行ログを残す

2つ目の使い道の例として、回帰の精度評価でMSE, RMSEの両方を算出する場合を考える。なお、MSEは二乗平均誤差、RMSEは二乗平均平方根誤差である。

ここで、MSE, RMSEはそれぞれ別の関数で算出するものとする。また、予測値と実測値は、リスト、NumPy配列、PandasのSeries, DataFrameのいずれの型で与えられても算出できるようにしたい。

このとき、予測値と実測値を1つの型に変換する処理をデコレータにまとめることで、MSEとRMSEを算出する関数はその型のみに対応した処理のみ記述すれば済むようになる。

ここではデコレータでNumPy配列に変換するものとして、例を以下に示す。

import numpy as np

def conv2nparray(func):
    def wrapper(*args):
        a = np.array(args[0]).flatten()
        b = np.array(args[1]).flatten()
        res = func(a, b)
        return res
    return wrapper

@conv2nparray
def mse(a, b):
    return np.mean((a-b)**2)

@conv2nparray
def rmse(a, b):
    return np.sqrt(np.mean((a-b)**2))

conv2nparrayでは、2つの引数をそれぞれ1次元のNumPy配列に変換している。また、配列の長さが等しくない場合の処理や、NaNを含むときの処理が必要であれば、デコレータの中に書くことも可能である。

参考

#

Pythonのデコレータについて - Qiita