Add decorators to make common use cases easy.

brian-brazil · brian-brazil · commit 20d8be02290c · 2015-02-10T00:22:02.000Z
diff --git a/prometheus_client/__init__.py b/prometheus_client/__init__.py
@@ -2,7 +2,10 @@
 
 import copy
 import re
+import time
+from contextlib import contextmanager
 from BaseHTTPServer import BaseHTTPRequestHandler
+from functools import wraps
 from threading import Lock
 
 __all__ = ['Counter', 'Gauge', 'Summary']
@@ -165,6 +168,23 @@ def inc(self, amount=1):
     with self._lock:
       self._value += amount
 
+  @contextmanager
+  def trackBlockRaises(self):
+    """Decorator to increment if a block raises an exception."""
+    try:
+      yield
+    except Exception, e:
+      self.inc()
+      raise e
+
+  def trackFunctionRaises(self, f):
+    """Decorator to increment if a function raises an exception."""
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+      with self.trackBlockRaises():
+        return f(*args, **kwargs)
+    return wrapper
+
   def _samples(self):
     with self._lock:
       return (('', {}, self._value), )
@@ -187,10 +207,31 @@ def dec(self, amount=1):
       self._value -= amount
 
   def set(self, value):
-    """Set gauge to the given amount."""
+    """Set gauge to the given value."""
     with self._lock:
       self._value = float(value)
 
+  def setToCurrentTime(self, value):
+    """Set gauge to the current unixtime."""
+    self.set(time.time())
+
+  @contextmanager
+  def trackBlockInprogress(self):
+    """Decorator to track how many of a block are in progress."""
+    self.inc()
+    try:
+      yield
+    finally:
+      self.dec()
+
+  def trackFunctionInprogress(self, f):
+    """Decorator to track how many of a function are in progress."""
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+      with self.trackBlockInprogress():
+        return f(*args, **kwargs)
+    return wrapper
+
   def _samples(self):
     with self._lock:
       return (('', {}, self._value), )
@@ -209,6 +250,24 @@ def observe(self, amount):
       self._count += 1
       self._sum += amount
 
+  @contextmanager
+  def timeBlock(self):
+    """Context manager to time how long a block of code takes in seconds."""
+    start = time.time()
+    try:
+      yield
+    finally:
+      # Time can go backwards.
+      self.observe(max(time.time() - start, 0))
+
+  def timeFunction(self, f):
+    """Decorator to time long a function takes in seconds."""
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+      with self.timeBlock():
+        return f(*args, **kwargs)
+    return wrapper
+
   def _samples(self):
     with self._lock:
       return (
@@ -257,7 +316,6 @@ def do_GET(self):
   s = Summary('ss', 'A summary', ['a', 'b'])
   s.labels('c', 'd').observe(17)
  
- 
   from BaseHTTPServer import HTTPServer
   server_address = ('', 8000)
   httpd = HTTPServer(server_address, MetricsHandler)
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -18,6 +18,34 @@ def test_increment(self):
   def test_negative_increment_raises(self):
     self.assertRaises(ValueError, self.counter.inc, -1)
 
+  def test_function_decorator(self):
+    @self.counter.trackFunctionRaises
+    def f(r):
+      if r:
+        raise Exception
+    f(False)
+    self.assertEquals(0, self.registry.get_sample_value('c'))
+    raised = False
+    try:
+      f(True)
+    except:
+      raised = True
+    self.assertTrue(raised)
+    self.assertEquals(1, self.registry.get_sample_value('c'))
+
+  def test_block_decorator(self):
+    with self.counter.trackBlockRaises():
+      pass
+    self.assertEquals(0, self.registry.get_sample_value('c'))
+    raised = False
+    try:
+      with self.counter.trackBlockRaises():
+        raise Exception
+    except:
+      raised = True
+    self.assertTrue(raised)
+    self.assertEquals(1, self.registry.get_sample_value('c'))
+
 class TestGauge(unittest.TestCase):
   def setUp(self):
     self.registry = CollectorRegistry()
@@ -32,6 +60,20 @@ def test_gauge(self):
     self.gauge.set(9)
     self.assertEquals(9, self.registry.get_sample_value('g'))
 
+  def test_function_decorator(self):
+    self.assertEquals(0, self.registry.get_sample_value('g'))
+    @self.gauge.trackFunctionInprogress
+    def f():
+      self.assertEquals(1, self.registry.get_sample_value('g'))
+    f()
+    self.assertEquals(0, self.registry.get_sample_value('g'))
+
+  def test_block_decorator(self):
+    self.assertEquals(0, self.registry.get_sample_value('g'))
+    with self.gauge.trackBlockInprogress():
+      self.assertEquals(1, self.registry.get_sample_value('g'))
+    self.assertEquals(0, self.registry.get_sample_value('g'))
+
 class TestSummary(unittest.TestCase):
   def setUp(self):
     self.registry = CollectorRegistry()
@@ -44,6 +86,20 @@ def test_summary(self):
     self.assertEquals(1, self.registry.get_sample_value('s_count'))
     self.assertEquals(10, self.registry.get_sample_value('s_sum'))
 
+  def test_function_decorator(self):
+    self.assertEquals(0, self.registry.get_sample_value('s_count'))
+    @self.summary.timeFunction
+    def f():
+      pass
+    f()
+    self.assertEquals(1, self.registry.get_sample_value('s_count'))
+
+  def test_block_decorator(self):
+    self.assertEquals(0, self.registry.get_sample_value('s_count'))
+    with self.summary.timeBlock():
+      pass
+    self.assertEquals(1, self.registry.get_sample_value('s_count'))
+
 class TestMetricWrapper(unittest.TestCase):
   def setUp(self):
     self.registry = CollectorRegistry()
