From d18381005a15e3caabc296651c4dc24808c1e9c1 Mon Sep 17 00:00:00 2001
From: Brian Brazil <brian.brazil@gmail.com>
Date: Mon, 6 Apr 2015 14:14:18 +0100
Subject: [PATCH] Add support for callbacks for gauges.

---
 README.md                     |  8 ++++++++
 prometheus_client/__init__.py | 12 ++++++++++++
 tests/test_client.py          |  9 +++++++++
 3 files changed, 29 insertions(+)

diff --git a/README.md b/README.md
index cc71faaf..35e36a65 100644
--- 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
index 7e2896e4..a0fd65c0 100644
--- 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
index 3b113cb9..4945131c 100644
--- 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):