Zenn
⌛

Pythonで関数の時間を計測するラッパーを作成

2022/12/07に公開
Python
#
wrapper
tech

結論

import datetime
from functools import wraps

def calc_func_time(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        start = datetime.datetime.now()
        result = f(*args, **kwargs)
        end = datetime.datetime.now()
        print(f'elapsed time: {f.__name__}: { end - start }')
        return result
    return wrapper

@calc_func_time
def func(*args, **kwargs):
    #(ある処理)
    pass

内容

ある処理(func())の時間を計測する最も古典的な方法は以下です。

import time

def func(*args, **kwargs):
    #(ある処理)
    pass
start = time.time()
func()
end = time.time()

print(f'time: {func.__name__}: { end - start }')

このプログラムを毎回書くのは面倒なので関数化したいです。
例えば以下のような関数を作って実行するとします。

def calc_func_time(func, *args, **kwargs):
    start = time.time()
    result = func(*args, **kwargs)
    end = time.time()
    print(f'time: {func.__name__}: { end - start }')
    return result

calc_func_time(func)

これでも一件落着かに思えますが、funcを呼び度にcalc_func_timeを書くのは面倒ですし、何よりmainでfuncが実行されているか一目でわかりにくいです。

そこでfuncに時間を計測する関数をラップしましょう。
あとプログラム完成前の余談ですが、time.time()の結果はformatしないとユーザーフレンドリーじゃないのでdatatime.datatime.now()を使いましょう。

以下が結論のプログラムです。

import datetime
from functools import wraps

def calc_func_time(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        start = datetime.datetime.now()
        result = f(*args, **kwargs)
        end = datetime.datetime.now()
        print(f'time: {f.__name__}: { end - start }')
        return result
    return wrapper

@calc_func_time
def func(*args, **kwargs):
    #(ある処理)
    pass
    
func()

これでfunc()を実行する度に時間を計測して出力してくれます。

毎回時間を計測させたくなかったら、do_printのような引数を用意した以下のようにラッパーを実装すればいいです。

def calc_func_time(f, do_print = True):
    @wraps(f)
    def wrapper(*args, **kwargs):
        start = datetime.datetime.now()
        result = f(*args, **kwargs)
        end = datetime.datetime.now()
	if do_print:
	    print(f'time: {f.__name__}: { end - start }')
        return result
    return wrapper

functools.wrapsについてはこちらで紹介しています。

https://zenn.dev/sergicalsix/articles/a1f5ff0b6a4a4c

Discussion