pythonで実行時間を測定しつつ時間も取得したい

実行時間を測定するには?

KaggleMasterのアライさん著・Kaggleコード遺産でおなじみ timer を使うのが便利です! いつもお世話になってます😊

Kaggleコード遺産 - Qiita

from contextlib import contextmanager
from time import time

# https://qiita.com/kaggle_master-arai-san/items/d59b2fb7142ec7e270a5#timer をちょっといじったやつ
@contextmanager
def timer(logger=None, format_str='{:.3f}[s]', prefix=None, suffix=None):
    if prefix: format_str = str(prefix) + format_str
    if suffix: format_str = format_str + str(suffix)
    start = time()
    yield
    d = time() - start
    out_str = format_str.format(d)
    if logger:
        logger.info(out_str)
    else:
        print(out_str)

あとは、測りたいコードの部分を with 区でくくります。かんたんで便利ですね!

with timer(prefix='foo'):
    # run some function
    sleep(10)

しかし timer は print で時間を表示してくれるものの時間の計測が timer 側で閉じてしまっているので、計測後に実行時間を取得することができません。たとえば実行時間を console に出しつつデータフレームなどに保存したい〜というときなど、実行時間自体にアクセスしたいですね。

Classを使って書き直す

これを解決する一つの方法が timer 自体に時間などを保存できるようにする方法です。timer ごとに値を属性として保存するため funcition だったものを Timer class に書き直しました

class Timer:
    def __init__(self, logger=None, format_str='{:.3f}[s]', prefix=None, suffix=None, sep=' '):

        if prefix: format_str = str(prefix) + sep + format_str
        if suffix: format_str = format_str + sep + str(suffix)
        self.format_str = format_str
        self.logger = logger
        self.start = None
        self.end = None

    @property
    def duration(self):
        if self.end is None:
            return 0
        return self.end - self.start

    def __enter__(self):
        self.start = time()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time()
        out_str = self.format_str.format(self.duration)
        if self.logger:
            self.logger.info(out_str)
        else:
            print(out_str)

enter / exit を定義して with のときの動作を書いているだけです。

これを使うと一度定義した timer に start_at / duration などの属性がひも付きますのであとから時間を取り出すことが可能です。

hoge_timer = Timer()

with hoge_timer:
    sleep(10)

print(hoge_timer.duration)

また今まで通り timer として使いたいよ〜というニーズもあると思います。その場合以下のように function を定義しておけば OK です

def timer(logger=None, format_str='{:.3f}[s]', prefix=None, suffix=None, sep=' '):
    return Timer(logger=logger, format_str=format_str, prefix=prefix, suffix=suffix, sep=sep)