Add the GaugeHistogram

brian-brazil · brian-brazil · commit 494a27d4fcab · 2018-09-07T17:17:03.000+01:00
Signed-off-by: Brian Brazil &lt;brian.brazil@robustperception.io&gt;
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -152,7 +152,7 @@ def get_sample_value(self, name, labels=None):
 '''The default registry.'''
 
 _METRIC_TYPES = ('counter', 'gauge', 'summary', 'histogram', 
-        'untyped', 'info', 'stateset')
+        'gaugehistogram', 'untyped', 'info', 'stateset')
 
 
 class Metric(object):
@@ -329,6 +329,33 @@ def add_metric(self, labels, buckets, sum_value):
         self.samples.append((self.name + '_sum', dict(zip(self._labelnames, labels)), sum_value))
 
 
+class GaugeHistogramMetricFamily(Metric):
+    '''A single gauge histogram and its samples.
+
+    For use by custom collectors.
+    '''
+    def __init__(self, name, documentation, buckets=None, labels=None):
+        Metric.__init__(self, name, documentation, 'gaugehistogram')
+        if labels is not None and buckets is not None:
+            raise ValueError('Can only specify at most one of buckets and labels.')
+        if labels is None:
+            labels = []
+        self._labelnames = tuple(labels)
+        if buckets is not None:
+            self.add_metric([], buckets)
+
+    def add_metric(self, labels, buckets):
+        '''Add a metric to the metric family.
+
+        Args:
+          labels: A list of label values
+          buckets: A list of pairs of bucket names and values.
+              The buckets must be sorted, and +Inf present.
+        '''
+        for bucket, value in buckets:
+            self.samples.append((self.name + '_bucket', dict(list(zip(self._labelnames, labels)) + [('le', bucket)]), value))
+
+
 class InfoMetricFamily(Metric):
     '''A single info and its samples.
 
diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
@@ -78,6 +78,10 @@ def generate_latest(registry=core.REGISTRY):
             mtype = 'gauge'
         elif mtype == 'stateset':
             mtype = 'gauge'
+        elif mtype == 'gaugehistogram':
+            # A gauge histogram is really a gauge,
+            # but this captures the strucutre better.
+            mtype = 'histogram'
         output.append('# HELP {0} {1}'.format(
             mname, metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
         output.append('\n# TYPE {0} {1}\n'.format(mname, mtype))
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -18,6 +18,7 @@
     GaugeMetricFamily,
     Histogram,
     HistogramMetricFamily,
+    GaugeHistogramMetricFamily,
     Info,
     InfoMetricFamily,
     Enum,
@@ -545,6 +546,18 @@ def test_histogram_labels(self):
         self.assertEqual(2, self.registry.get_sample_value('h_count', {'a': 'b'}))
         self.assertEqual(3, self.registry.get_sample_value('h_sum', {'a': 'b'}))
 
+    def test_gaugehistogram(self):
+        self.custom_collector(GaugeHistogramMetricFamily('h', 'help', buckets=[('0', 1), ('+Inf', 2)]))
+        self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '0'}))
+        self.assertEqual(2, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
+
+    def test_gaugehistogram_labels(self):
+        cmf = GaugeHistogramMetricFamily('h', 'help', labels=['a'])
+        cmf.add_metric(['b'], buckets=[('0', 1), ('+Inf', 2)])
+        self.custom_collector(cmf)
+        self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'a': 'b', 'le': '0'}))
+        self.assertEqual(2, self.registry.get_sample_value('h_bucket', {'a': 'b', 'le': '+Inf'}))
+
     def test_info(self):
         self.custom_collector(InfoMetricFamily('i', 'help', value={'a': 'b'}))
         self.assertEqual(1, self.registry.get_sample_value('i_info', {'a': 'b'}))
diff --git a/tests/test_exposition.py b/tests/test_exposition.py
@@ -13,6 +13,7 @@
 from prometheus_client import CollectorRegistry, generate_latest
 from prometheus_client import push_to_gateway, pushadd_to_gateway, delete_from_gateway
 from prometheus_client import CONTENT_TYPE_LATEST, instance_ip_grouping_key
+from prometheus_client.core import GaugeHistogramMetricFamily
 from prometheus_client.exposition import default_handler, basic_auth_handler, MetricsHandler
 
 try:
@@ -28,6 +29,12 @@ class TestGenerateText(unittest.TestCase):
     def setUp(self):
         self.registry = CollectorRegistry()
 
+    def custom_collector(self, metric_family):
+        class CustomCollector(object):
+            def collect(self):
+                return [metric_family]
+        self.registry.register(CustomCollector())
+
     def test_counter(self):
         c = Counter('cc', 'A counter', registry=self.registry)
         c.inc()
@@ -72,6 +79,10 @@ def test_histogram(self):
 hh_sum 0.05
 ''', generate_latest(self.registry))
 
+    def test_gaugehistogram(self):
+        self.custom_collector(GaugeHistogramMetricFamily('gh', 'help', buckets=[('1.0', 4), ('+Inf', (5))]))
+        self.assertEqual(b'''# HELP gh help\n# TYPE gh histogram\ngh_bucket{le="1.0"} 4.0\ngh_bucket{le="+Inf"} 5.0\n''', generate_latest(self.registry))
+
     def test_info(self):
         i = Info('ii', 'A info', ['a', 'b'], registry=self.registry)
         i.labels('c', 'd').info({'foo': 'bar'})
