Combine the context manager and the decorator.

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):
         """

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()