from collections.abc import Iterator
from contextlib import contextmanager
from time import perf_counter


class TimingLog:
    def __init__(self) -> None:
        self._steps: dict[str, float] = {}

    @contextmanager
    def measure(self, name: str) -> Iterator[None]:
        started = perf_counter()
        try:
            yield
        finally:
            self._steps[name] = round(perf_counter() - started, 3)

    def to_dict(self) -> dict[str, float]:
        total = round(sum(self._steps.values()), 3)
        return {**self._steps, "total": total}