コンテキストマネージャーを使った処理時間の計測

この記事は Pythonのコードを短く簡潔に書くテクニック Advent Calendar 2017 の18日目です。

はじめに

Pythonのスクリプト内で特定部分の処理時間を計測したいときに、以下のようなコードを書いたことがある人は多いと思います。

start = time.time()

# 計測したい処理
# :

end = time.time()
print(end - start)

@contextlib.contextmanagerデコレーターを使ってコンテキストマネージャーを作りwith文で呼び出せば、これをもっと簡潔で再利用可能なコードに置き換えることができます。

with simple_timer('処理1'):
    # 計測したい処理
    # :

例題1

以下の様に２つある処理を別々に計測したい場合を考えてみます。
処理内容はダミーとして指定時間スループするだけにします。

# 処理1
time.sleep(0.1)

# 処理2
time.sleep(0.2)

普通の書き方

start = time.time()

# 処理1
time.sleep(0.1)

end = time.time()
print('処理1: {:.3f}'.format(end-start))

start = time.time()

# 処理2
time.sleep(0.2)

end = time.time()
print('処理2: {:.3f}'.format(end-start))

出力

処理1: 0.105
処理2: 0.202

時間計測のコードが混ざるので元々やっていた処理がわかりにくくなってしまいます。

コンテキストマネージャーで書いた場合

コンテキストマネージャーを使うことで、時間計測のコードを計測対象のコードから分離して定義できます。
先ずは@contextlib.contextmanagerデコレーターを使ってコンテキストマネージャーを定義します。

from contextlib import contextmanager

@contextmanager
def simple_timer(label):
    start = time.time()
    yield
    end = time.time()
    print('{}: {:.3f}'.format(label, end-start))

あとは計測したい範囲をwith文のブロックで囲めば(インデントすれば)処理時間が計測できます。

with simple_timer('処理1'):
    time.sleep(0.1)

with simple_timer('処理2'):
    time.sleep(0.2)

出力

処理1: 0.103
処理2: 0.203

例題2

次は、先ほどのコードが繰り返し呼ばれる場合に、２つの処理を別々に集計して計測したい場合を考えてみます。

for _ in range(10):
    # 処理1
    time.sleep(0.1)
    # 処理2
    time.sleep(0.2)

普通に書いた場合

times = defaultdict(float)

for _ in range(10):
    start = time.time()

    # 処理1
    time.sleep(0.1)

    end = time.time()
    times['処理1'] += end - start

    start = time.time()

    # 処理2
    time.sleep(0.2)

    end = time.time()
    times['処理2'] += end - start

for label, t in times.items():
    print('{}: {:.3f}'.format(label, t))

処理1: 1.035
処理2: 2.020

コンテキストマネージャーで書いた場合

コンテキストマネージャーではyield文で呼び出し元に値を渡すことができます。
ここで時間計測用のコンテキストマネージャーを渡すようにします。

@contextmanager
def sum_timer():

    times = defaultdict(float)

    @contextmanager
    def timer(label):
        start = time.time()
        yield
        end = time.time()
        times[label] += end - start

    yield timer

    for label, t in times.items():
        print('{}: {:.3f}'.format(label, t))

呼び出し側ではwith文の後ろのasでyieldから渡された値を受け取ることができます。

with sum_timer() as timer:

    for _ in range(10):

        with timer('処理1'):
            time.sleep(0.1)

        with timer('処理2'):
            time.sleep(0.2)

処理1: 1.021
処理2: 2.021

例題3

今度は更に全体の処理時間も出力するようにしてみます。

普通に書いた場合

面倒なので省略

コンテキストマネージャーで書いた場合

コンテキストマネージャー側に全体の処理時間を計測する処理を追加するだけです。
呼び出し側のコードには影響しません。

@contextmanager
def total_timer(total_label):

    times = defaultdict(float)

    @contextmanager
    def timer(label):
        start = time.time()
        yield
        end = time.time()
        times[label] += end - start

    with timer(total_label):
        yield timer

    for label, t in times.items():
        if label != total_label:
            print('{}: {:.3f}'.format(label, t))
    print('{}: {:.3f}'.format(total_label, times[total_label]))

with total_timer('全体') as timer:

    for _ in range(10):

        with timer('処理1'):
            time.sleep(0.1)

        with timer('処理2'):
            time.sleep(0.2)

出力

処理1: 1.014
処理2: 2.036
全体: 3.051

他にも、呼び出し回数・平均・最小・最大を出力するといった応用も可能です。

参考

• Python3ドキュメント
  • ライブラリーリファレンス
    • 1. 組み込み型
      • 4.11. コンテキストマネージャ型
    • 29.6. contextlib — with 文コンテキスト用ユーティリティ
      • @contextlib.contextmanager
  • 言語リファレンス
    • 1. データモデル
      • 3.3.8. with文とコンテキストマネージャ
    • 1. 複合文 (compound statement)
      • 8.5. with 文