Merge pull request prometheus#24 from brian-brazil/gauge-callback

brian-brazil · brian-brazil · commit 6b9e6dee8df0 · 2015-04-06T14:30:54.000+01:00
Add support for callbacks for gauges.
diff --git a/README.md b/README.md
@@ -113,6 +113,14 @@ with g.track_inprogress():
   pass
 ```
 
+A Gauge can also take it's value from a callback:
+
+```python
+d = Gauge('data_objects', 'Number of objects')
+my_dict = {}
+d.set_function(lambda: len(my_dict))
+```
+
 ### Summary
 
 Summaries track the size and number of events.
diff --git a/prometheus_client/__init__.py b/prometheus_client/__init__.py
@@ -8,6 +8,7 @@
 import os
 import time
 import threading
+import types
 try:
     from BaseHTTPServer import BaseHTTPRequestHandler
     from BaseHTTPServer import HTTPServer
@@ -279,6 +280,17 @@ def wrapped(*args, **kwargs):
 
         return InprogressTracker(self)
 
+    def set_function(self, f):
+        '''Call the provided function to return the Gauge value.
+
+        The function must return a float, and may be called from
+        multiple threads.
+        All other methods of the Gauge become NOOPs.
+        '''
+        def samples(self):
+            return (('', {}, float(f())), )
+        self._samples = types.MethodType(samples, self)
+
     def _samples(self):
         with self._lock:
             return (('', {}, self._value), )
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -87,6 +87,15 @@ def test_block_decorator(self):
             self.assertEqual(1, self.registry.get_sample_value('g'))
         self.assertEqual(0, self.registry.get_sample_value('g'))
 
+    def test_gauge_function(self):
+        x = {}
+        self.gauge.set_function(lambda: len(x))
+        self.assertEqual(0, self.registry.get_sample_value('g'))
+        self.gauge.inc()
+        self.assertEqual(0, self.registry.get_sample_value('g'))
+        x['a'] = None
+        self.assertEqual(1, self.registry.get_sample_value('g'))
+
 
 class TestSummary(unittest.TestCase):
     def setUp(self):
