timed.py

"""Time any function"""

from asyncio import iscoroutinefunction
from collections.abc import Awaitable, Callable, Coroutine
from contextlib import contextmanager
from functools import wraps
from inspect import isawaitable
from logging import StreamHandler, getLogger
from logging.handlers import RotatingFileHandler
from sys import stdout
from time import perf_counter_ns
from typing import Any, ParamSpec, TypeVar, overload

logger = getLogger("timings")

logger.addHandler(StreamHandler(stdout))
logger.addHandler(RotatingFileHandler("timings.log", maxBytes=10_000))

logger.setLevel("INFO")

P = ParamSpec("P")
R = TypeVar("R")


@contextmanager
def timer(infos: str):
    t0 = perf_counter_ns()
    yield
    t1 = perf_counter_ns()
    logger.info(
        "timing - %s - %s ms",
        infos,
        round((t1 - t0) / 1_000_000),
    )


@overload
def timed(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]: ...


@overload
def timed(func: Callable[P, R]) -> Callable[P, R]: ...


def timed(func: Callable[P, R]) -> Callable[P, R] | Callable[P, Coroutine[Any, Any, R]]:
    """Log execution time of a function. Support both sync and async function and is
    fully typed.
    """
    if iscoroutinefunction(func):

        @wraps(func)
        async def with_timing(*args: P.args, **kwargs: P.kwargs) -> R:
            with timer(func.__module__ + "." + func.__name__):
                return await func(*args, **kwargs)

    else:

        @wraps(func)
        def with_timing(*args: P.args, **kwargs: P.kwargs) -> R:
            with timer(func.__module__ + "." + func.__name__):
                return func(*args, **kwargs)

    return with_timing