mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-06 08:16:00 +00:00
53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
import time
|
|
|
|
from loguru import logger
|
|
|
|
|
|
class Timer(object):
|
|
def __init__(self, name: str, use_ms: bool = False, stack_level: int = 2):
|
|
self.name: str = name
|
|
self.use_ms: bool = use_ms
|
|
self.stack_level: int = stack_level
|
|
|
|
self.time_start: float = 0
|
|
self.time_end: float = 0
|
|
self.time_cost: float = 0
|
|
|
|
def __enter__(self, *args, **kwargs):
|
|
self.time_start = time.time()
|
|
logger.info(f"---------- enter {self.name} ----------", stacklevel=self.stack_level)
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
self.time_end = time.time()
|
|
self.time_cost = self.time_end - self.time_start
|
|
if self.use_ms:
|
|
time_str = f"{self.time_cost * 1000:.2f}ms"
|
|
else:
|
|
time_str = f"{self.time_cost:.3f}s"
|
|
|
|
logger.info(f"---------- leave {self.name} [{time_str}] ----------", stacklevel=self.stack_level)
|
|
|
|
|
|
def timer(name: str = None, use_ms: bool = False, stack_level: int = 2):
|
|
def decorator(func):
|
|
def wrapper(*args, **kwargs):
|
|
with Timer(name=name or func.__name__, use_ms=use_ms, stack_level=stack_level + 1):
|
|
return func(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import random
|
|
|
|
|
|
@timer("run_func_final", use_ms=True)
|
|
def run_func():
|
|
time.sleep(random.uniform(0.05, 0.15))
|
|
print("done")
|
|
|
|
|
|
run_func()
|