Python: context manager to measure elapsed time.

measure_time.py

"""
This script will print:

INFO:root:test
INFO:root:elapsed: 3.003285470069386 (or something similar)
"""

from collections.abc import Callable, Iterator
from contextlib import contextmanager
import time

@contextmanager
def measure_time(callback: Callable[[float], None]) -> Iterator[None]:
    begin = time.perf_counter()
    try:
        yield
    finally:
        callback(time.perf_counter() - begin)


import logging
logging.basicConfig(level=logging.INFO)

with measure_time(lambda t: logging.info(f"elapsed: {t}")):
    time.sleep(3)
    logging.info("test")