Add UntypedMetricFamily to prometheus_client.core (prometheus#220)

idsvandermolen · brian-brazil · commit 9828994d106b · 2017-12-08T12:23:40.000Z
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -176,6 +176,29 @@ def __eq__(self, other):
                 and self.samples == other.samples)
 
 
+class UntypedMetricFamily(Metric):
+    '''A single untyped metric and its samples.
+    For use by custom collectors.
+    '''
+    def __init__(self, name, documentation, value=None, labels=None):
+        Metric.__init__(self, name, documentation, 'untyped')
+        if labels is not None and value is not None:
+            raise ValueError('Can only specify at most one of value and labels.')
+        if labels is None:
+            labels = []
+        self._labelnames = tuple(labels)
+        if value is not None:
+            self.add_metric([], value)
+
+    def add_metric(self, labels, value):
+        '''Add a metric to the metric family.
+        Args:
+        labels: A list of label values
+        value: The value of the metric.
+        '''
+        self.samples.append((self.name, dict(zip(self._labelnames, labels)), value))
+
+
 class CounterMetricFamily(Metric):
     '''A single counter and its samples.
 
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -321,6 +321,16 @@ def collect(self):
                 return [metric_family]
         self.registry.register(CustomCollector())
 
+    def test_untyped(self):
+        self.custom_collector(UntypedMetricFamily('u', 'help', value=1))
+        self.assertEqual(1, self.registry.get_sample_value('u', {}))
+
+    def test_untyped_labels(self):
+        cmf = UntypedMetricFamily('u', 'help', labels=['a', 'c'])
+        cmf.add_metric(['b', 'd'], 2)
+        self.custom_collector(cmf)
+        self.assertEqual(2, self.registry.get_sample_value('u', {'a': 'b', 'c': 'd'}))
+
     def test_counter(self):
         self.custom_collector(CounterMetricFamily('c', 'help', value=1))
         self.assertEqual(1, self.registry.get_sample_value('c', {}))
@@ -370,6 +380,9 @@ def test_histogram_labels(self):
         self.assertEqual(3, self.registry.get_sample_value('h_sum', {'a': 'b'}))
 
     def test_bad_constructors(self):
+        self.assertRaises(ValueError, UntypedMetricFamily, 'u', 'help', value=1, labels=[])
+        self.assertRaises(ValueError, UntypedMetricFamily, 'u', 'help', value=1, labels=['a'])
+
         self.assertRaises(ValueError, CounterMetricFamily, 'c', 'help', value=1, labels=[])
         self.assertRaises(ValueError, CounterMetricFamily, 'c', 'help', value=1, labels=['a'])
 
@@ -390,6 +403,8 @@ def test_bad_constructors(self):
         self.assertRaises(KeyError, HistogramMetricFamily, 'h', 'help', buckets={}, sum_value=1)
 
     def test_labelnames(self):
+        cmf = UntypedMetricFamily('u', 'help', labels=iter(['a']))
+        self.assertEqual(('a',), cmf._labelnames)
         cmf = CounterMetricFamily('c', 'help', labels=iter(['a']))
         self.assertEqual(('a',), cmf._labelnames)
         gmf = GaugeMetricFamily('g', 'help', labels=iter(['a']))
