memoryscope.core.utils.timer 源代码

import time
from typing import Literal

from memoryscope.core.utils.logger import Logger

TIME_LOG_TYPE = Literal["end", "wrap", "none"]


[文档] class Timer(object): """ A class used to measure the execution time of code blocks. It supports logging the elapsed time and can be customized to display time in seconds or milliseconds. """
[文档] def __init__(self, name: str, time_log_type: TIME_LOG_TYPE = "end", use_ms: bool = True, stack_level: int = 2, float_precision: int = 4, **kwargs): """ Initializes the `Timer` instance with the provided args and sets up a logger Args: name (str): The log name. time_log_type (str): The log type. Defaults to 'End'. use_ms (bool): Use 'ms' as the timescale or not. Defaults to True. stack_level (int): The stack level of log. Defaults to 2. float_precision (int): The precision of cost time. Defaults to 4. """ self.name: str = name self.time_log_type: TIME_LOG_TYPE = time_log_type self.use_ms: bool = use_ms self.stack_level: int = stack_level self.float_precision: int = float_precision self.kwargs: dict = kwargs # time recorder self.t_start = 0 self.t_end = 0 self.cost = 0 self.logger = Logger.get_logger()
def _set_cost(self): """ Accumulate the cost time. """ self.t_end = time.time() self.cost = self.t_end - self.t_start if self.use_ms: self.cost *= 1000 @property def cost_str(self): """ Represent the cost time into a formatted string. """ self._set_cost() if self.use_ms: return f"cost={self.cost:.4f}ms" else: return f"cost={self.cost:.4f}s" def __enter__(self, *args, **kwargs): """ Begin timing. """ self.t_start = time.time() if self.time_log_type == "wrap": self.logger.info(f"----- {self.name}.begin -----") return self def __exit__(self, exc_type, exc_value, exc_tb): """ End timing and print the formatted log. """ if self.time_log_type == "none": return lines = [] if self.time_log_type == "wrap": lines.append(f"----- {self.name}.end -----") else: lines.append(self.name) lines.append(self.cost_str) if self.kwargs: for k, v in self.kwargs.items(): if isinstance(v, float): float_style = f".{self.float_precision}f" line = f"{k}={v:{float_style}}" else: line = f"{k}={v}" lines.append(line) self.logger.info(" ".join(lines), stacklevel=self.stack_level)
[文档] def timer(func): """ A decorator function that measures the execution time of the wrapped function. Args: func (Callable): The function to be wrapped and timed. Returns: Callable: The wrapper function that includes timing functionality. """ def wrapper(*args, **kwargs): """ The wrapper function that manages the timing of the original function. Args: *args: Variable length argument list for the decorated function. **kwargs: Arbitrary keyword arguments for the decorated function. Returns: Any: The result of the decorated function. """ with Timer(name=func.__name__, **kwargs): return func(*args, **kwargs) return wrapper