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

この記事は [Pythonのコードを短く簡潔に書くテクニック Advent Calendar 2017](https://qiita.com/advent-calendar/2017/python-short-code) の18日目です。

# はじめに

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

```py3
start = time.time()

# 計測したい処理
# :

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

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

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

[@contextlib.contextmanager]:https://docs.python.jp/3/library/contextlib.html#contextlib.contextmanager
[コンテキストマネージャー]:https://docs.python.jp/3/library/stdtypes.html#typecontextmanager
[with文]:http://docs.python.jp/3/reference/compound_stmts.html#the-with-statement

# 例題1

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

```py3
# 処理1
time.sleep(0.1)

# 処理2
time.sleep(0.2)
```

## 普通の書き方

```py3
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]デコレーターを使ってコンテキストマネージャーを定義します。

```py3
from contextlib import contextmanager

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

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

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

with simple_timer('処理2'):
    time.sleep(0.2)
```
```:出力
処理1: 0.103
処理2: 0.203
```

# 例題2

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

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

## 普通に書いた場合

```py3
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文で呼び出し元に値を渡すことができます。
ここで時間計測用のコンテキストマネージャーを渡すようにします。

```py3:
@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`から渡された値を受け取ることができます。

```py3
with sum_timer() as timer:

    for _ in range(10):

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

        with timer('処理2'):
            time.sleep(0.2)
```
```py3
処理1: 1.021
処理2: 2.021
```

# 例題3

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

## 普通に書いた場合

面倒なので省略

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

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

```py3
@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]))
```
```py3
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ドキュメント
    - ライブラリーリファレンス
        - 4. 組み込み型
            - [4.11. コンテキストマネージャ型](https://docs.python.jp/3/library/stdtypes.html#typecontextmanager)
        - 29.6. contextlib — with 文コンテキスト用ユーティリティ
            - [@contextlib.contextmanager](https://docs.python.jp/3/library/contextlib.html#contextlib.contextmanager)
    - 言語リファレンス
        - 3. データモデル
            - [3.3.8. with文とコンテキストマネージャ](https://docs.python.jp/3/reference/datamodel.html#with-statement-context-managers)
        - 8. 複合文 (compound statement)
            - [8.5. with 文](https://docs.python.jp/3/reference/compound_stmts.html#with)