Scoped timer `with` statement python

timer.py

import math
import threading
import time
import types


def format_duration(seconds: float, sigfig: int = 3) -> str:
    if seconds == 0:
        return "0s"

    units = [
        (1, "s"),
        (1e-3, "ms"),
        (1e-6, "μs"),
        (1e-9, "ns"),
    ]

    for factor, unit in units:
        if seconds >= factor:
            value = seconds / factor
            digits = sigfig - int(math.floor(math.log10(abs(value)))) - 1
            formatted_value = round(value, digits)
            return f"{formatted_value:g}{unit}"

    value = seconds / 1e-9
    digits = sigfig - int(math.floor(math.log10(abs(value)))) - 1
    formatted_value = round(value, digits)
    return f"{formatted_value:g}ns"


class Timer:
    start: float
    label: str
    _local = threading.local()

    COLORS = [
        "\033[96m",  # Cyan
        "\033[93m",  # Yellow
        "\033[95m",  # Magenta
    ]
    RESET = "\033[0m"

    def __init__(self, label: str):
        self.label = label

    def __enter__(self):
        self.start = time.perf_counter()
        if not hasattr(self._local, "timing_stack"):
            self._local.timing_stack = []
        self._local.timing_stack.append(self)

        depth = len(self._local.timing_stack) - 1
        color = self.COLORS[depth % len(self.COLORS)]

        indent = ""
        for i in range(depth):
            level_color = self.COLORS[i % len(self.COLORS)]
            indent += f"{level_color}│{self.RESET}"

        print(f"{indent}{color}╭ {self.label}{self.RESET}")
        return self

    def __exit__(
        self,
        _exc_type: type[BaseException] | None,
        _exc_value: BaseException | None,
        _traceback: types.TracebackType | None,
    ) -> bool | None:
        end = time.perf_counter()
        elapsed = end - self.start

        depth = len(self._local.timing_stack) - 1
        color = self.COLORS[depth % len(self.COLORS)]

        indent = ""
        for i in range(depth):
            level_color = self.COLORS[i % len(self.COLORS)]
            indent += f"{level_color}│{self.RESET}"

        print(f"{indent}{color}╰ {self.label} [{format_duration(elapsed)}]{self.RESET}")
        self._local.timing_stack.pop()