Add ability to time durations for gauges.

brian-brazil · brian-brazil · commit e25541aa6a38 · 2015-08-19T13:27:52.000+01:00
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -285,6 +285,32 @@ def wrapped(*args, **kwargs):
 
         return InprogressTracker(self)
 
+    def time(self):
+        '''Time a block of code or function, and set the duration in seconds.
+
+        Can be used as a function decorator or context manager.
+        '''
+
+        class Timer(object):
+            def __init__(self, gauge):
+                self._gauge = gauge
+
+            def __enter__(self):
+                self._start = time.time()
+
+            def __exit__(self, typ, value, traceback):
+                # Time can go backwards.
+                self._gauge.set(max(time.time() - self._start, 0))
+
+            def __call__(self, f):
+                @wraps(f)
+                def wrapped(*args, **kwargs):
+                    with self:
+                        return f(*args, **kwargs)
+                return wrapped
+
+        return Timer(self)
+
     def set_function(self, f):
         '''Call the provided function to return the Gauge value.
 
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -1,6 +1,7 @@
 from __future__ import unicode_literals
 import os
 import threading
+import time
 import unittest
 
 
@@ -108,6 +109,22 @@ def test_gauge_function(self):
         x['a'] = None
         self.assertEqual(1, self.registry.get_sample_value('g'))
 
+    def test_function_decorator(self):
+        self.assertEqual(0, self.registry.get_sample_value('g'))
+
+        @self.gauge.time()
+        def f():
+            time.sleep(.001)
+
+        f()
+        self.assertNotEqual(0, self.registry.get_sample_value('g'))
+
+    def test_block_decorator(self):
+        self.assertEqual(0, self.registry.get_sample_value('g'))
+        with self.gauge.time():
+            time.sleep(.001)
+        self.assertNotEqual(0, self.registry.get_sample_value('g'))
+
 
 class TestSummary(unittest.TestCase):
     def setUp(self):
