src/common/timer.py

"""Timer class.

Context and decorator form timer.
"""

import contextlib
from functools import wraps
from time import perf_counter

from src.common.logger import log_success, log_info


class Timer(contextlib.ContextDecorator):
    """Timer.

    Examples:
        >>> with Timer('Code1'):
        ...     sleep(1)
        * Code1        | 1.00s (0.02m)
    """

    def __init__(self, name="Elapsed time"):
        self.name = name

    def __enter__(self):
        log_info(f"{'[START] ' + self.name:15}")
        self.start_time = perf_counter()
        return self

    def __exit__(self, *exc):
        elapsed_time = perf_counter() - self.start_time
        log_success(f"{self.name:15} | {elapsed_time:.2f}s ({elapsed_time/60:.2f}m)")
        return False


def T(fn: callable) -> callable:
    """Timer decorator.

    Example:
        >>> @T
        >>> def f():
        ...     sleep(1)
        * Elapsed time | 1.00s (0.02m)
    """

    @wraps(fn)
    def _log(*args, **kwargs):
        with Timer(fn.__name__):
            rst = fn(*args, **kwargs)
        return rst

    return _log