File size: 1,081 Bytes
615dee0 0157dfd 615dee0 0157dfd 615dee0 0157dfd 615dee0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import time
from typing import Dict, List
from dotenv import load_dotenv
import os
from src.utils.log import logger
load_dotenv()
ENABLE_TIMING = os.getenv("ENABLE_TIMING", "false").lower() == "true"
if ENABLE_TIMING:
logger.info("Timing is enabled")
TimingType = Dict[str, List[float]]
TIMES: TimingType = {}
def timeit(name: str):
def _timeit(func):
def wrapper(*args, **kwargs):
start = time.time()
result = func(*args, **kwargs)
end = time.time()
if name not in TIMES:
TIMES[name] = []
TIMES[name].append(end - start)
return result
if ENABLE_TIMING:
return wrapper
return func
return _timeit
def get_times() -> TimingType:
_warn_if_timing_disabled()
return TIMES
def reset_times() -> None:
_warn_if_timing_disabled()
TIMES.clear()
def _warn_if_timing_disabled() -> None:
if not ENABLE_TIMING:
logger.warning(
"Timing is disabled, please set ENABLE_TIMING to true in the .env file")
|