Pythonで関数・メソッドの実行時間を一行で計測するデコレータを作る

やりたいこと

Pythonで関数やメソッドの実行時間を計算したい。Jupyter Notebookならマジックコマンド%%timeitによってセルごとの実行時間を記録することができるが、通常のPythonスクリプトで関数に一行付け足すだけでprintしてくれるようなものが欲しい。

実現方法

関数の実行前後に時刻を計測して引き算すれば良い。これはデコレータの考え方とフィットするため、デコレータとして以下のように実装する。

ポイントは以下の通り。

• 通常の関数でもメソッドでも同様に使える
• 計測結果は秒単位で出力するが、目的に応じて小数点以下何桁まで表示するかをndigitsという引数で指定できる
• 複数の関数で実行した時に見分けられるよう、messageという引数を与えられるようにした
  • ネストしているのはデコレータに引数を与えるため
• デコレータを使うとエディタなどで元の関数のdocstringが隠蔽されてしまうことがあるため、functools.wrapsを使用した

import time
from functools import wraps

def timeit(message, ndigits=2):
    """Print execution time [sec] of function/method
    - message: message to print with time
    - ndigits: precision after the decimal point
    """
    def outer_wrapper(func):
        # @wraps: keep docstring of "func"
        @wraps(func)
        def inner_wrapper(*args, **kwargs):
            start = time.time()
            result = func(*args, **kwargs)
            end = time.time()
            print(" 🥑 Excecution time for '{message}': {sec} [sec]".format(
                message=message,
                sec=round(end-start, ndigits))
            )
            return result
        return inner_wrapper
    return outer_wrapper

使い方

与えられた秒数だけsleepする関数sleeper()の実行時間を計測したいとする。以下のように関数定義の前に先ほど定義したデコレータを記述する。

@timeit("sleeper", ndigits=2)
def sleeper(sec):
    time.sleep(sec)
    return sec

この場合、sleeper()を実行するたびに以下のようなテキストがprintされる。

 🥑 Excecution time for 'sleeper': 1.24 [sec]

個人的にはこれで十分だが、print以外の機能を追加したり、時間の単位をmsなどに変更できるようにしても良いと思う。

以上。