Pythonのコンテキストの仕組みを使って処理時間を良い感じに表示する

プログラムの処理時間を計測したいようなケースがたまにあります。処理が簡単な場合は

import time

def some_func():
    ...
    
b = time.perf_counter()
some_func()
print(f'{time_perf_counter() - b} secs')

とかで十分なのですが

def foo():
    bar()
    
def bar():
    baz1()
    baz2()
    
def baz1():
    ....
    
def baz2():
    ....

のようになってくると（そして実際にはこんなケースの方が多い）、foo の前後で時間を計測しても実際に遅いのがどこなのか特定できません。

そこで

def foo():
    with measure('call bar'):
        bar()
    
def bar():
    with measure('call baz1'):
        baz1()
	
    with measure('call baz2'):
        baz2()
    
def baz1():
    ....
    
def baz2():
    ....
    
with measure('call foo') as m:
    foo()
    
print_measure(m)

[call foo]: 1.232 secs
    [call bar]: 1.211 secs
        [call baz1]: 0.01082 secs
        [call baz2]: 1.2 secs

なような事が実現できればfooの呼び出しで支配的なのはbaz2の実行であることがわかります。ここで、いきなり登場したmeasureとprint_measureは自前で実装したコードですが、以下の通りで60行程度のものです。ちなみに以下のコードはSelf型を利用しているのでPython3.11以降じゃないと文法エラーになります(Self型便利）。

import time
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Self


class Measure:
    def __init__(self, name: str) -> None:
        self._name = name
        self._elapsed: float = 0
        self._children: list[Self] = []

    @property
    def name(self) -> str:
        return self._name

    @property
    def elapsed(self) -> float:
        return self._elapsed

    def set_elapsed(self, value: float):
        self._elapsed = value

    def add_child(self, m: Self):
        self._children.append(m)


_context: ContextVar[Measure | None] = ContextVar('measure', default=None)


@contextmanager
def measure(name: str):
    parent = _context.get()
    m = Measure(name)
    token = _context.set(m)
    b = time.perf_counter()

    yield m

    a = time.perf_counter()
    _context.reset(token)
    m.set_elapsed(a - b)

    if parent:
        parent.add_child(m)


def print_measure(m: Measure, indent_level: int = 4):
    def _print_measure(m: Measure, indent: int):
        spaces = ' ' * indent
        print(f'{spaces}[{m.name}]: {m.elapsed:.4} secs')

        for c in m._children:
            _print_measure(c, indent + indent_level)

    _print_measure(m, 0)

Measureは所謂Composite patternでmeasureのネストに対応したものになっていて、あるwith measureの呼び出しの内部で呼び出されるwith measureの結果はchildrenに格納されます。contextvarsモジュールはPython3.7から導入されたもので、Reactを触っている方であれば、React.Contextと同様な機構と思ってもらって問題ないと思います。
with measureの実行ごとにコンテキストに格納してある状態をすげ替える事でMeasureの構造を構築していっています。
print_measureの方はMeasureの構造を深さ優先で探索して結果を表示しているだけです。

ちなみにシングルスレッドを仮定できればwith measureの呼び出しごとに状態をスタックに積んでいけばよいので、以下のようなコードでも同じ事は実現できます。

_stack: list[Measure] = []


@contextmanager
def measure(name: str):
    parent = _stack[-1] if _stack else None
    m = Measure(name)
    _stack.append(m)
    b = time.perf_counter()

    yield m

    a = time.perf_counter()
    _stack.pop()
    m.set_elapsed(a - b)

    if parent:
        parent.add_child(m)

コンテキスト変数はスレッドごとに独立しているので、実際はわざわざ上記のようにスタックを管理する必要はなく素直にコンテキスト変数を使っておけばよいでしょう。

これくらいのコード量で結構便利な仕組みが作れるので、コンテキストは便利ですね。