Python小ネタ: withとcontextmanagerで実行時間を計測する

きっかけ

Kaggleで実行時間に制限のあるコンペのsolutionを眺めていたときに見つけた記法を、後で見返せるようにメモしておきます。
https://www.youtube.com/watch?v=msUWJPXWCUM

実行時間を計測する

コードサンプル

import time

@contextmanager
def timer(name):
    t0 = time.time()
    yield
    print(f'[{name}] done in {time.time() - t0:.0f} s')

with timer('process train'):
    (処理)

with

with 文は、ブロックの実行を、コンテキストマネージャによって定義されたメソッドでラップするために使われます

ということは、個人的にはあまり意識してきませんでしたが、これまでよく使ってきた

try:
    f = open('workfile')
    read_data = f.read()
finally:
    f.close()

↓

with open('workfile') as f:
    read_data = f.read()

で出てくるopenはコンテキストマネージャということになります。

コンテキストマネージャ(context manager)

コンテキストマネージャ(context manager) とは、 with 文の実行時にランタイムコンテキストを定義するオブジェクトです。コンテキストマネージャは、コードブロックを実行するために必要な入り口および出口の処理を扱います。コンテキストマネージャは通常、 with 文（ with 文 の章を参照）により起動されますが、これらのメソッドを直接呼び出すことで起動することもできます。

例を拝借すると、
http://python.keicode.com/lang/context-manager.php

class MyContextManager:

    def __enter__(self):
        print 'ENTER!'
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print 'EXIT!'

    def foo(self):
        print 'FOO!'

# TEST PROGRAM

with MyContextManager() as cm:
    print '** hello **'
    cm.foo()

の実行結果がこのようになります。

ENTER!
** hello **
FOO!
EXIT!

with ブロックで行う処理の前後に確実に実行したい処理を、__enter()__メソッドと_exit()_メソッドにそれぞれ記述しておくという使い方ができます。

@contextlib.contextmanager デコレータ

この関数は with 文コンテキストマネージャのファクトリ関数を定義するために利用できる デコレータ です。

このデコレータを使うと、コンテキストマネージャをより簡単に実装することができます。
別の例を拝借すると、

from contextlib import contextmanager

@contextmanager
def tag(name):
    print("<%s>" % name)
    yield
    print("</%s>" % name)

with tag("h1"):
   print("foo")

の実行結果がこのようになります。（引用元にあるように"実際に HTML を生成する方法としてはお勧めできません!"）

<h1>
foo
</h1>

tag関数がジェネレータ（return の代わりに yield で値を返す）となり、@contextlib.contextmanager デコレータによってコンテキストマネージャになっています。
yieldの前後にあるprint文が、それぞれ __enter__() と __exit__() に対応しています。

参考

https://docs.python.jp/3/library/contextlib.html
http://blog.amedama.jp/entry/2015/10/02/234946
http://python.keicode.com/lang/context-manager.php