context.py (1922B)
1 #Embedded file name: C:\BuildAgent\work\20d73f1b24eed06e\eve\release\V24.01\packages\datadog\dogstatsd\context.py 2 from functools import wraps 3 from time import time 4 from datadog.util.compat import is_higher_py35, iscoroutinefunction 5 if is_higher_py35(): 6 from datadog.dogstatsd.context_async import _get_wrapped_co 7 else: 8 9 def _get_wrapped_co(self, func): 10 raise NotImplementedError(u'Decorator `timed` compatibility with coroutine functions requires Python 3.5 or higher.') 11 12 13 class TimedContextManagerDecorator(object): 14 15 def __init__(self, statsd, metric = None, tags = None, sample_rate = 1, use_ms = None): 16 self.statsd = statsd 17 self.metric = metric 18 self.tags = tags 19 self.sample_rate = sample_rate 20 self.use_ms = use_ms 21 self.elapsed = None 22 23 def __call__(self, func): 24 if not self.metric: 25 self.metric = '%s.%s' % (func.__module__, func.__name__) 26 if iscoroutinefunction(func): 27 return _get_wrapped_co(self, func) 28 29 @wraps(func) 30 def wrapped(*args, **kwargs): 31 start = time() 32 try: 33 return func(*args, **kwargs) 34 finally: 35 self._send(start) 36 37 return wrapped 38 39 def __enter__(self): 40 if not self.metric: 41 raise TypeError('Cannot used timed without a metric!') 42 self.start = time() 43 return self 44 45 def __exit__(self, type, value, traceback): 46 self._send(self.start) 47 48 def _send(self, start): 49 elapsed = time() - start 50 use_ms = self.use_ms if self.use_ms is not None else self.statsd.use_ms 51 elapsed = int(round(1000 * elapsed)) if use_ms else elapsed 52 self.statsd.timing(self.metric, elapsed, self.tags, self.sample_rate) 53 self.elapsed = elapsed 54 55 def start(self): 56 self.__enter__() 57 58 def stop(self): 59 self.__exit__(None, None, None)