Add a context manager and decorator for timing metrics.

README.md

@@ -37,6 +37,19 @@ client.metric('some_metric', {'value_a': 100, 'value_b': 100, 'value_c': True})
 
 # Records a single value with one tag
 client.metric('some_metric', 123, tags={'server_name': 'my-server'})
+
+# Times a function call. The measurement name defaults to __module__.__func__
+@client.timer('some_metric', tags={'server_name':'my-server'})
+def some_operation():
+    pass
+
+# Uses a context manager to time an operation
+with client.timer('some_metric'):
+    time.sleep(2)
+
+# As above, but time in milliseconds
+with client.timer('some_metric', use_ms=True):
+    time.sleep(2)
 ```
 
 #### Global tags

setup.py

@@ -2,7 +2,6 @@
 
 from setuptools import find_packages, setup
 
-
 here = os.path.abspath(os.path.dirname(__file__))
 
 about = {}

telegraf/__init__.py

@@ -1,6 +1,6 @@
 from __future__ import absolute_import
 
-from .__version__ import __title__, __description__, __version__  # noqa
-from .client import TelegrafClient, HttpClient
+from .__version__ import __description__, __title__, __version__  # noqa
+from .client import HttpClient, TelegrafClient
 
 __all__ = ('TelegrafClient', 'HttpClient')

telegraf/client.py

@@ -1,6 +1,20 @@
+import socket
 from abc import abstractmethod
+from functools import wraps
+from time import time
+
 from telegraf.protocol import Line
-import socket
+from telegraf.utils import is_higher_py35
+
+if is_higher_py35():
+    from telegraf.context import wrapped_coroutine
+    from asyncio import iscoroutinefunction
+else:
+    def wrapped_coroutine(self, func):
+        raise NotImplementedError(u"Async timer decorator requires Python 3.5 or higher.")
+
+    def iscoroutinefunction(*args, **kwargs):
+        return False
 
 
 class ClientBase(object):
@@ -28,6 +42,9 @@ def metric(self, measurement_name, values, tags=None, timestamp=None):
         line = Line(measurement_name, values, all_tags, timestamp)
         self.send(line.to_line_protocol())
 
+    def timer(self, measurement_name, tags=None, use_ms=False):
+        return TimerHelper(self, measurement_name=measurement_name, tags=tags, use_ms=use_ms)
+
     @abstractmethod
     def send(self, data):
         pass
@@ -77,3 +94,50 @@ def send(self, data):
         this issues the request in the background.
         """
         self.future_session.post(url=self.url, data=data)
+
+
+class TimerHelper(object):
+    def __init__(self, client, measurement_name=None, tags=None, use_ms=False):
+        self.client = client
+        self.measurement_name = measurement_name
+        self.tags = tags or {}
+        self.use_ms = use_ms
+
+    def __call__(self, func):
+        """Decorator helper for timing function calls."""
+        if not self.measurement_name:
+            self.measurement_name = '%s.%s' % (func.__module__, func.__name__)
+
+        # Coroutines
+        if iscoroutinefunction(func):
+            return wrapped_coroutine(self, func)
+
+        @wraps(func)
+        def wrapped(*args, **kwargs):
+            start = time()
+            try:
+                return func(*args, **kwargs)
+            finally:
+                self._send(start)
+        return wrapped
+
+    def __enter__(self):
+        if not self.measurement_name:
+            raise TypeError("No metric name specified.")
+        self.start_time = time()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        # Report the elapsed time of the context manager.
+        self._send(self.start_time)
+
+    def _send(self, start_time):
+        elapsed = time() - start_time
+        tags = self.tags.copy()
+        if self.use_ms:
+            elapsed = int(round(1000 * elapsed))
+            tags['units'] = 'ms'
+        else:
+            tags['units'] = 's'
+        line = Line(self.measurement_name, elapsed, tags, None)
+        self.client.send(line.to_line_protocol())

telegraf/context.py

@@ -0,0 +1,20 @@
+"""
+Decorator for async methods.
+
+Requires Python 3.5 or higher.
+"""
+# stdlib
+from functools import wraps
+from time import time
+
+
+def wrapped_coroutine(self, func):
+    """Timing wrapper for async functions."""
+    @wraps(func)
+    async def wrapped_co(*args, **kwargs):  # noqa
+        start = time()
+        try:
+            return await func(*args, **kwargs)
+        finally:
+            self._send(start)
+    return wrapped_co

telegraf/defaults/django.py

@@ -1,10 +1,9 @@
 from __future__ import absolute_import
-from django.conf import settings
 
+from django.conf import settings
 from telegraf import defaults
 from telegraf.client import TelegrafClient
 
-
 telegraf = None
 
 if telegraf is None:

telegraf/tests.py

@@ -1,10 +1,11 @@
 # -*- coding: utf-8 -*-
 
-from telegraf.client import ClientBase, TelegrafClient, HttpClient
-from telegraf.protocol import Line
-from telegraf.utils import format_string, format_value
 import unittest
+
 import mock
+from telegraf.client import ClientBase, HttpClient, TelegrafClient
+from telegraf.protocol import Line
+from telegraf.utils import format_string, format_value
 
 
 class TestLine(unittest.TestCase):

telegraf/utils.py

@@ -1,3 +1,5 @@
+import sys
+
 try:
     basestring
 except NameError:
@@ -46,3 +48,8 @@ def format_value(value):
     elif isinstance(value, float):
         value = str(value)
     return value
+
+
+def is_higher_py35():
+    """Check that the Python is version 3.5 or higher."""
+    return sys.version_info >= (3, 5)