contextmanagerを使ってIPythonの%%timeitのようなものをつくりたい

Question

contextmanagerを使ってwithブロック内の処理を複数回実行し平均の速度を計る機能を考えています。
JupyterNotebooksなどのIPython系での%%timeitの簡易版と思っていただければと思います。

下記のコードを書いてみたのですが、AttributeError: __enter__が出てしまいます。
contextmanager_timer()の部分をreturn contextmanager_timer()とするとエラーは出ないのですが、1度しか実行されないためにコードの意味がなくなります。
どう書き換えればよいでしょうか?

また、計測する処理にprint()などが含まれていた場合os.devnullに流して静かに実行することも考えているのですが、そちらについてもアドバイスがあれば伺いたいです。

[Wandbox]三へ( へ՞ਊ ՞)へ ﾊｯﾊｯ

import time
from contextlib import contextmanager


def simple_timeit(label: str, exec_limit: int = 1000, time_limit: int = 5):
    @contextmanager
    def contextmanager_timer():
        start = time.time()
        yield
        nonlocal total_time
        total_time += time.time() - start

    cnt = 0
    total_time = 0
    while cnt < exec_limit and total_time < time_limit:
        cnt += 1
        contextmanager_timer()
    print(f"{label}: {total_time / cnt:.8f}s on average of {cnt} runs.")


def main():
    with simple_timeit("for loops"):
        for _ in range(10 ** 7):
            pass


if __name__ == "__main__":
    main()

for loops: 0.00000000s on average of 1000 runs.
Traceback (most recent call last):
  File "prog.py", line 28, in <module>
    main()
  File "prog.py", line 22, in main
    with simple_timeit("for loops"):
AttributeError: __enter__

Answer 1

余談です。（このデコレータ自体は参考にはしないでください

with my_timeit を for _ in my_timeit に置き換えるデコレータ

@recompile_with_timeit
def main():
    with my_timeit(repeat=3):
        print("TEST")


if __name__ == '__main__':
    main()

Output:

start
TEST
TEST
TEST
end

デコレーターが必要になる時点で本末転倒感がありますが、興味本位で実装してみました。
(コードは後述)

※ 対象のコードから、目的の動作を得られるように強引に実装しただけなので、
決して質問に対する回答としてのお勧めの解決策ではありません(重要)。デメリットが多数 あります。

• pylint 等でエラーになる
• デバッガのステップ実行でも対応するコードが得られない

もし with -> for _ in のタイプ量が問題なら、スニペット登録する方が賢い解決方法です。

興味深い点があって、こちらが本題
IPython の timeit でも、使い方は異なりますが
実行時にコードから関数を作り、何か構文木(AST)を操作するような事をしてました。

IPythonより引用

        # This codestring is taken from timeit.template - we fill it in as an
        # AST, so that we can apply our AST transformations to the user code
        # without affecting the timing code.
        timeit_ast_template = ast.parse('def inner(_it, _timer):\n'
                                        '    setup\n'
                                        '    _t0 = _timer()\n'
                                        '    for _i in _it:\n'
                                        '        stmt\n'
                                        '    _t1 = _timer()\n'
                                        '    return _t1 - _t0\n')

        timeit_ast = TimeitTemplateFiller(ast_setup, ast_stmt).visit(timeit_ast_template)
        timeit_ast = ast.fix_missing_locations(timeit_ast)

def inner(_it, _timer):
    setup
    _t0 = _timer()
    for _i in _it:
        stmt
    _t1 = _timer()
    return _t1 - _t0

という関数の構文木を生成して、setup と stmt をコードに置き換える。
推測ですが、stmt 置き換えは、大量に繰り返す場合の関数呼び出しのオーバーヘッド削減でしょうか。

時間計測の処理を最適化するのに、timeitの実装は参考になると思いました。
因みに、標準モジュールのtimeitでも、文字列ベースで同様の事を行ってます。

---

デコレータの実装

#!/usr/bin/env python3.8
"""
rewrite with-generator AST
"""

import ast
import inspect


def my_timeit(repeat):
    print("start")
    yield from range(repeat)
    print("end")


class Transformer(ast.NodeTransformer):
    """
    AST変換

    - 関数の再定義にデコレータが2重に呼ばれてしまうので、削除
    - with my_timeit -> for _ in my_timeit に書き換え

    """
    def visit_FunctionDef(self, node):
        self.generic_visit(node)
        node.decorator_list = [
            x for x in node.decorator_list if x.id != "recompile_with_timeit"
        ]
        return node

    def visit_With(self, node):
        expr = node.items[0].context_expr
        if isinstance(expr, ast.Call) and expr.func.id == "my_timeit":
            return ast.For(
                target=ast.Name(id="_", ctx=ast.Store()),
                iter=expr,
                body=node.body,
                orelse=[])
        return node


def recompile_with_timeit(func):
    src = inspect.getsource(func)

    node = ast.parse(src)
    node = Transformer().visit(node)
    node = ast.fix_missing_locations(node)

    new_code = compile(node, "<string>", "exec")
    namespace = {}
    exec(new_code, globals(), namespace)
    return namespace.get(func.__name__)


@recompile_with_timeit
def main():
    "Test code"

    with my_timeit(repeat=10):
        print("TEST")

    for _ in range(10):
        pass

    # 他の with ステートメントに影響がないことを確認
    with open(__file__, encoding="utf-8") as stream:
        for line in map(str.rstrip, stream):
            print(line)


if __name__ == '__main__':
    main()

Comments

[退会済みユーザー]

ちょうど最近ASTの勉強をしていたので、こういう風でも使えるのかーと驚いています。
recompile_with_timeitを理解できるように、処理の流れを丹念に追いかけてみます。

Answer 2

contextmanagerを使ってwithブロック内の処理を複数回実行し

with ではできません。

同様の事が実現できる言語(Smalltalk等)では、
ブロック・オブジェクトを引数に取ることができますが、
Pythonでは code オブジェクト、関数オブジェクト辺りがこれに相当します。

問題点は、コンテキストマネージャからwith のブロックを再利用する手段がない事。
with では yield により実行コンテキストが「切り替わる」といった方が解りやすいかもしれません。

テストとして実行することを考えると、
デコレーター + 関数として実装が率直な解決策です。
※ 実装は略

@timeit_dec(repeat=10)
def test_for_loop(): # <- デコレータに渡るのはここの関数オブジェクト
    print("TEST")

一時的なケースで使いたいなら代案
記述は with -> for _ in になりますが、
ジェネレーターとして実装する方法 ... こちらの方がやりたいことに近いのではないでしょうか。

def my_timeit(repeat):
    print("start")

    yield from range(repeat)

    print("end")


if __name__ == '__main__':
    for _ in my_timeit(repeat=10):
        print("TEST")

ここから更に前処理・後処理の部分は、
デコレーターやコンテキストマネージャーに括り出せます。

---

計測する処理にprint()などが含まれていた場合os.devnullに流して静かに実行すること

• Python では sys.stdout にファイルライクなオブジェクトを指定するとリダイレクトできます。
• 局所的なリダイレクトは contextlib.redirect_stdout も併せて調べて見て下さい。
• print のみなら

__builtins__.print 自体を差し替えると、若干内部のメソッド呼び出しが減らせるかもしれません。
が、logging モジュールを使っている場合等も考慮すると、標準入力のリダイレクトに分があり。

但し、マルチスレッドの場合等は、単純にリダイレクトだけとはいきません、
テストは別プロセスで実行する等の工夫が必要です。

IOの速度を測るのでなければ、基本的に単体テストが必要なコードからはprint文は省けるはずです。
可能ならIOは切り離して関数を実装した方がテストしやすくなります。

---

問題点: 時間制限を実装しようとしてますが
以下のコードはループが完了するまで、時間チェックに処理が戻らない為
意図した時間で終わりません。

for _ in range(10 ** 7):
    pass

経過時間チェック

繰り返し毎に毎回チェックされるようにする必要があります。

for _ in range(10 * 7):
     経過時間のチェック
     pass

→ テスト対象のコードを直ぐに終わるような小さなコードに分割する
while 文で単純な分割が難しい場合は、ジェネレーターにするなどの工夫が必要です。

import time

def my_timeit(label, repeat, timeout=float("inf"), count=0):
    start = time.time()
    time_limit = start + timeout

    for _ in range(repeat):
        if time.time() > time_limit:
            break
        count += 1
        yield
        
    total_time = time.time() - start
    
    print(f"{label}: {total_time / count:.8f}s on average of {count} runs.")


if __name__ == '__main__':
    for _ in my_timeit("test", repeat=20, timeout=10):
        print("TEST")
        time.sleep(1)

追記: 処理内容によっては誤差かもしれないけど、
「経過時間チェック」により時間は遅くなるので、用途次第ではこの実装は適切でないかもしれません。
他の案: 一定時間経過後にシグナルを送る等、テストコードに影響しない他の仕組みで解決。
参考: pytest-timeout

追記2: yymmt さんの投稿を見て、気づいた点
time.time は time.perf_timer の方が適切でした。

time.time()はシステムに設定された時刻に依存するので

この関数が返す値は通常減少していくことはありませんが、この関数を 2 回呼び出し、その呼び出しの間にシステムクロックの時刻を巻き戻して設定した場合には、以前の呼び出しよりも低い値が返ることがあります。

timeitの実装も timer.perf_counter を利用しています。

Answer 3

少し考えてみましたがwithブロックの中をループする処理は出来ないように思います。contextmanagerは単に__enter__と__exit__を使いやすくしたものですので

with loop_func():
    do_something()

という構文は

__enter__で定義した前処理
do_something()
__exit__で定義した後処理

として展開されます。この形式でループされるようには記述出来ないためです。そこで外で回すか、中で回すか、という話になりますが、これをコードにして表すと

# 外で回す
for _ in range(n):
    __enter__で定義した前処理
    do_something()
    __exit__で定義した後処理

# 以下と等価な処理
# for _ in range(n):
#     with simple_timeit():
#         do_something()

もしくは

# 中で回す
__enter__で定義した前処理
for _ in range(n):
    do_something()
__exit__で定義した後処理

# 以下と等価な処理
# with simple_timeit():
#     for _ in range(n):
#         do_something()

となり、両者ともwith構文でシンプルに記述するという目的とは異なっていると思います。with構文を諦めて単純なforループにするのであれば対応方法はありますが、やはりこれもやりたい事とは違うのだろうなぁと思います。参考までにサンプルコードを示します。

import io
import sys
import time
from contextlib import redirect_stdout, redirect_stderr


def simple_timeit(
    label: str,
    exec_limit: int = 1000,
    time_limit: int = 5,
    stdout: io.TextIOWrapper = sys.stdout,
    stderr: io.TextIOWrapper = sys.stderr,
):
    with redirect_stdout(stdout), redirect_stderr(stderr):
        cnt = 0
        total_time = 0.0
        while cnt < exec_limit and total_time < time_limit:
            start = time.perf_counter()
            yield cnt
            total_time += time.perf_counter() - start
            cnt += 1
    print(f"{label}: {total_time / cnt:.8f}s on average of {cnt} runs.")
    return cnt


def main():
    for i in simple_timeit("for loops", stdout=None, stderr=None):
        print(i)
        print("stdout: hello")
        print("stderr: world", file=sys.stderr)


if __name__ == "__main__":
    main()