import functools from time import perf_counter as time_now def safe_wraps(wrapper, *args, **kwargs): """Safely wraps partial functions.""" while isinstance(wrapper, functools.partial): wrapper = wrapper.func return functools.wraps(wrapper, *args, **kwargs) class Timer: """A context manager/decorator for statsd.timing().""" def __init__(self, client, stat, rate=1, tags=None, use_decimals=True, auto_send=True): self.client = client self.stat = stat self.rate = rate self.tags = tags self.ms = None self._sent = False self._start_time = None self.use_decimals = use_decimals self.auto_send = auto_send def __call__(self, f): """Thread-safe timing function decorator.""" @safe_wraps(f) def _wrapped(*args, **kwargs): start_time = time_now() try: return f(*args, **kwargs) finally: elapsed_time_ms = 1000.0 * (time_now() - start_time) self.client.timing(self.stat, elapsed_time_ms, self.rate, self.tags, self.use_decimals) self._sent = True return _wrapped def __enter__(self): return self.start() def __exit__(self, typ, value, tb): self.stop(send=self.auto_send) def start(self): self.ms = None self._sent = False self._start_time = time_now() return self def stop(self, send=True): if self._start_time is None: raise RuntimeError('Timer has not started.') dt = time_now() - self._start_time self.ms = 1000.0 * dt # Convert to milliseconds. if send: self.send() return self def send(self): if self.ms is None: raise RuntimeError('No data recorded.') if self._sent: raise RuntimeError('Already sent data.') self._sent = True self.client.timing(self.stat, self.ms, self.rate, self.tags, self.use_decimals)