|
22 | 22 | # logger.add(sys.stdout, format=LOG_FORMAT, level="DEBUG", colorize=True) |
23 | 23 | import logging |
24 | 24 | import logging.config |
| 25 | +import time |
| 26 | +from collections.abc import Callable, Coroutine |
| 27 | +from functools import wraps |
25 | 28 | from os import getenv |
| 29 | +from typing import Any, ParamSpec, TypeVar |
26 | 30 |
|
27 | 31 | LOGGER_NAME = None |
28 | 32 |
|
@@ -93,3 +97,50 @@ def set_log_level_to_info(): |
93 | 97 |
|
94 | 98 | # Set initial log level |
95 | 99 | set_log_level_to_info() |
| 100 | + |
| 101 | + |
| 102 | +# Define generic type variables for return type and parameters |
| 103 | +R = TypeVar("R") |
| 104 | +P = ParamSpec("P") |
| 105 | + |
| 106 | + |
| 107 | +def time_execution_sync( |
| 108 | + additional_text: str = "", |
| 109 | +) -> Callable[[Callable[P, R]], Callable[P, R]]: |
| 110 | + def decorator(func: Callable[P, R]) -> Callable[P, R]: |
| 111 | + @wraps(func) |
| 112 | + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: |
| 113 | + start_time = time.time() |
| 114 | + result = func(*args, **kwargs) |
| 115 | + execution_time = time.time() - start_time |
| 116 | + logger.debug( |
| 117 | + f"{additional_text} Execution time: {execution_time:.2f} seconds" |
| 118 | + ) |
| 119 | + return result |
| 120 | + |
| 121 | + return wrapper |
| 122 | + |
| 123 | + return decorator |
| 124 | + |
| 125 | + |
| 126 | +def time_execution_async( |
| 127 | + additional_text: str = "", |
| 128 | +) -> Callable[ |
| 129 | + [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]] |
| 130 | +]: |
| 131 | + def decorator( |
| 132 | + func: Callable[P, Coroutine[Any, Any, R]] |
| 133 | + ) -> Callable[P, Coroutine[Any, Any, R]]: |
| 134 | + @wraps(func) |
| 135 | + async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: |
| 136 | + start_time = time.time() |
| 137 | + result = await func(*args, **kwargs) |
| 138 | + execution_time = time.time() - start_time |
| 139 | + logger.debug( |
| 140 | + f"{additional_text} Execution time: {execution_time:.2f} seconds" |
| 141 | + ) |
| 142 | + return result |
| 143 | + |
| 144 | + return wrapper |
| 145 | + |
| 146 | + return decorator |
0 commit comments