diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py index 1e13eabbc..70381247b 100644 --- a/datadog/dogstatsd/base.py +++ b/datadog/dogstatsd/base.py @@ -127,45 +127,31 @@ def timing(self, metric, value, tags=None, sample_rate=1): """ self._report(metric, 'ms', value, tags, sample_rate) - def timed(self, metric, tags=None, sample_rate=1): + class _TimedContextManagerDecorator(object): + """ + A context manager and a decorator which will report the elapsed time in + the context OR in a function call. """ - A decorator that will measure the distribution of a function's run - time. Optionally specify a list of tags or a sample rate. - :: - @statsd.timed('user.query.time', sample_rate=0.5) - def get_user(user_id): - # Do what you need to ... - pass + def __init__(self, statsd, metric, tags=None, sample_rate=1): + self.statsd = statsd + self.metric = metric + self.tags = tags + self.sample_rate = sample_rate + + def __call__(self, func): + """Decorator which returns the elapsed time of the function call.""" + manager = self - # Is equivalent to ... - start = time.time() - try: - get_user(user_id) - finally: - statsd.timing('user.query.time', time.time() - start) - """ - def wrapper(func): @wraps(func) def wrapped(*args, **kwargs): start = time() result = func(*args, **kwargs) - self.timing(metric, time() - start, tags=tags, - sample_rate=sample_rate) + manager.statsd.timing( + manager.metric, time() - start, tags=manager.tags, + sample_rate=manager.sample_rate) return result return wrapped - return wrapper - - class _TimedContextManager(object): - def __init__(self, statsd, metric, tags=None, sample_rate=1): - """ - Create a context manager which will report the time elapsed by its - context. - """ - self.statsd = statsd - self.metric = metric - self.tags = tags - self.sample_rate = sample_rate def __enter__(self): self.start = time() @@ -175,13 +161,20 @@ def __exit__(self, type, value, traceback): self.statsd.timing(self.metric, time() - self.start, self.tags, self.sample_rate) - def timed_context(self, metric, tags=None, sample_rate=1): + def timed(self, metric, tags=None, sample_rate=1): """ - A context manager that will measure the distribution of a context's run - time. Optionally specify a list of tags or a sample rate. + A decorator or context manager that will measure the distribution of a + function's/context's run time. Optionally specify a list of tags or a + sample rate. :: - with statsd.timed_context('user.query.time', sample_rate=0.5): + @statsd.timed('user.query.time', sample_rate=0.5) + def get_user(user_id): + # Do what you need to ... + pass + + # Is equivalent to ... + with statsd.timed('user.query.time', sample_rate=0.5): # Do what you need to ... pass @@ -192,7 +185,7 @@ def timed_context(self, metric, tags=None, sample_rate=1): finally: statsd.timing('user.query.time', time.time() - start) """ - return self._TimedContextManager(self, metric, tags, sample_rate) + return self._TimedContextManagerDecorator(self, metric, tags, sample_rate) def set(self, metric, value, tags=None, sample_rate=1): """ diff --git a/tests/unit/dogstatsd/test_statsd.py b/tests/unit/dogstatsd/test_statsd.py index d38c2ac5f..57954b55c 100644 --- a/tests/unit/dogstatsd/test_statsd.py +++ b/tests/unit/dogstatsd/test_statsd.py @@ -170,7 +170,7 @@ def func(a, b, c=1, d=1): self.assert_almost_equal(0.5, float(value), 0.1) def test_timed_context(self): - with self.statsd.timed_context('timed_context.test'): + with self.statsd.timed('timed_context.test'): time.sleep(0.5) packet = self.recv() @@ -187,7 +187,7 @@ class ContextException(Exception): pass def func(self): - with self.statsd.timed_context('timed_context.test.exception'): + with self.statsd.timed('timed_context.test.exception'): time.sleep(0.5) raise ContextException()