62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
import time
|
|
import logging
|
|
import asyncio
|
|
|
|
from functools import wraps
|
|
|
|
logger = logging.getLogger(__file__.split("/")[-1])
|
|
|
|
|
|
__LIMIT_TIMEIT_PRINT_MAX_LEN__ = 128
|
|
|
|
|
|
def timeit(
|
|
func=None, *, level: int = logging.DEBUG, output_args: bool = False, output_maxlen: int = __LIMIT_TIMEIT_PRINT_MAX_LEN__
|
|
):
|
|
if func is not None:
|
|
return timeit(level=level, output_args=output_args, output_maxlen=output_maxlen)(func)
|
|
|
|
def _get_timeit_arg_print(*args, **kwargs):
|
|
args_print = f"{args}"
|
|
args_print = args_print if len(args_print) < output_maxlen else args_print[: output_maxlen - 4] + "...)"
|
|
kwargs_print = f"{kwargs}"
|
|
kwargs_print = kwargs_print if len(kwargs_print) < output_maxlen else kwargs_print[: output_maxlen - 4] + "...}"
|
|
return args_print, kwargs_print
|
|
|
|
def _decorator(func):
|
|
if logger.level > level:
|
|
return func
|
|
is_async = asyncio.iscoroutinefunction(func)
|
|
|
|
@wraps(func)
|
|
async def _async_timeit_wrapper(*args, **kwargs):
|
|
if output_args:
|
|
args_print, kwargs_print = _get_timeit_arg_print(args, kwargs)
|
|
logger.log(level, f"async call {func.__name__} {args_print} {kwargs_print}")
|
|
else:
|
|
logger.log(level, f"async call {func.__name__}")
|
|
start_time = time.perf_counter()
|
|
result = await func(*args, **kwargs)
|
|
end_time = time.perf_counter()
|
|
total_time = end_time - start_time
|
|
logger.log(level, f"async quit {func.__name__} Took {total_time:.3f} seconds")
|
|
return result
|
|
|
|
@wraps(func)
|
|
def _sync_timeit_wrapper(*args, **kwargs):
|
|
if output_args:
|
|
args_print, kwargs_print = _get_timeit_arg_print(args, kwargs)
|
|
logger.log(level, f"sync call {func.__name__} {args_print} {kwargs_print}")
|
|
else:
|
|
logger.log(level, f"sync call {func.__name__}")
|
|
start_time = time.perf_counter()
|
|
result = func(*args, **kwargs)
|
|
end_time = time.perf_counter()
|
|
total_time = end_time - start_time
|
|
logger.log(level, f"sync quit {func.__name__} Took {total_time:.3f} seconds")
|
|
return result
|
|
|
|
return _async_timeit_wrapper if is_async else _sync_timeit_wrapper
|
|
|
|
return _decorator
|