profiling code for python inspired by clojure tufte library

profile.py

import time
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from contextlib import contextmanager
from collections import defaultdict
import threading

_profiling_data = threading.local()

@contextmanager
def profile(label: str):
    if not hasattr(_profiling_data, "timing"):
        _profiling_data.timing = defaultdict(float)

    start_time = time.perf_counter()
    try:
        yield
    finally:
        end_time = time.perf_counter()
        _profiling_data.timing[label] += end_time - start_time


@contextmanager
def profiling(label: str):
    if not hasattr(_profiling_data, "timing"):
        _profiling_data.timing = defaultdict(float)

    try:
        with profile('total'):
            yield
    finally:
        total_time = _profiling_data.timing['total']
        max_label_length = max((len(label) for label in _profiling_data.timing.keys()), default=5)
        label_width = max(max_label_length, 5)

        print(f"Profiling Summary for {label}:")
        print(f"{'Label':<{label_width}} {'Time (s)':>12} {'% Time':>10}")
        print("-" * (label_width + 24))
        for label, total_time_label in sorted(_profiling_data.timing.items(), key=lambda x: x[1], reverse=True):
            percentage = (total_time_label / total_time) * 100 if total_time > 0 else 0
            print(f"{label:<{label_width}} {total_time_label:>12.6f} {percentage:>9.2f}%")

        _profiling_data.timing.clear()


class ProfilingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        with profiling(f"Request: {request.method} {request.url.path}"):
            response = await call_next(request)
        return response