Expand unit support.

If unit is missing add it rather than erroring.
Add unit support to core types.
Don't allow unit on info/stateset/enum, as that doesn't make sense.

Signed-off-by: Brian Brazil <brian.brazil@robustperception.io>

prometheus_client/core.py

@@ -198,12 +198,12 @@ class Metric(object):
     and SummaryMetricFamily instead.
     '''
     def __init__(self, name, documentation, typ, unit=''):
+        if unit and not name.endswith("_" + unit):
+            name += "_" + unit
         if not _METRIC_NAME_RE.match(name):
             raise ValueError('Invalid metric name: ' + name)
         self.name = name
         self.documentation = documentation
-        if unit and not name.endswith("_" + unit):
-            raise ValueError("Metric name not suffixed by unit: " + name)
         self.unit = unit
         if typ == 'untyped':
             typ = 'unknown'
@@ -239,8 +239,8 @@ class UnknownMetricFamily(Metric):
     '''A single unknwon metric and its samples.
     For use by custom collectors.
     '''
-    def __init__(self, name, documentation, value=None, labels=None):
-        Metric.__init__(self, name, documentation, 'unknown')
+    def __init__(self, name, documentation, value=None, labels=None, unit=''):
+        Metric.__init__(self, name, documentation, 'unknown', unit)
         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:
@@ -265,11 +265,11 @@ class CounterMetricFamily(Metric):
 
     For use by custom collectors.
     '''
-    def __init__(self, name, documentation, value=None, labels=None, created=None):
+    def __init__(self, name, documentation, value=None, labels=None, created=None, unit=''):
         # Glue code for pre-OpenMetrics metrics.
         if name.endswith('_total'):
            name = name[:-6]
-        Metric.__init__(self, name, documentation, 'counter')
+        Metric.__init__(self, name, documentation, 'counter', unit)
         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:
@@ -735,14 +735,19 @@ def _samples(self):
 
 def _MetricWrapper(cls):
     '''Provides common functionality for metrics.'''
-    def init(name, documentation, labelnames=(), namespace='', subsystem='', registry=REGISTRY, **kwargs):
+    def init(name, documentation, labelnames=(), namespace='', subsystem='', unit='', registry=REGISTRY, **kwargs):
         full_name = ''
         if namespace:
             full_name += namespace + '_'
         if subsystem:
             full_name += subsystem + '_'
         full_name += name
 
+        if unit and not full_name.endswith("_" + unit):
+            full_name += "_" + unit
+        if unit and cls._type in ('info', 'stateset'):
+            raise ValueError('Metric name is of a type that cannot have a unit: ' + full_name)
+
         if cls._type == 'counter' and full_name.endswith('_total'):
             full_name = full_name[:-6]  # Munge to OpenMetrics.
 
@@ -767,7 +772,7 @@ def describe():
         collector.describe = describe
 
         def collect():
-            metric = Metric(full_name, documentation, cls._type)
+            metric = Metric(full_name, documentation, cls._type, unit)
             for suffix, labels, value in collector._samples():
                 metric.add_sample(full_name + suffix, labels, value)
             return [metric]

prometheus_client/openmetrics/parser.py

@@ -244,11 +244,15 @@ def build_metric(name, documentation, typ, unit, samples):
         if name in seen_metrics:
             raise ValueError("Duplicate metric: " + name)
         seen_metrics.add(name)
+        if unit and not name.endswith("_" + unit):
+            raise ValueError("Unit does not match metric name: " + name)
+        if unit and typ in ['info', 'stateset']:
+            raise ValueError("Units not allowed for this metric type: " + name)
         metric = core.Metric(name, documentation, typ, unit)
         # TODO: check labelvalues are valid utf8
         # TODO: check only histogram buckets have exemplars.
-        # TODO: Info and stateset can't have units
         # TODO: check samples are appropriately grouped and ordered
+        # TODO: check info/stateset values are 1/0
         # TODO: check for metadata in middle of samples
         # TODO: Check histogram bucket rules being followed
         metric.samples = samples

tests/openmetrics/test_parser.py

@@ -384,6 +384,8 @@ def test_invalid_input(self):
                 ('# UNIT a\t\n# EOF\n'),
                 ('# UNIT a seconds\n# EOF\n'),
                 ('# UNIT a_seconds seconds \n# EOF\n'),
+                ('# TYPE x_u info\n# UNIT x_u u\n# EOF\n'),
+                ('# TYPE x_u stateset\n# UNIT x_u u\n# EOF\n'),
                 # Bad metric names.
                 ('0a 1\n# EOF\n'),
                 ('a.b 1\n# EOF\n'),

tests/test_core.py

@@ -466,11 +466,22 @@ def test_invalid_names_raise(self):
         self.assertRaises(ValueError, Counter, 'c_total', '', labelnames=['a:b'])
         self.assertRaises(ValueError, Counter, 'c_total', '', labelnames=['__reserved'])
         self.assertRaises(ValueError, Summary, 'c_total', '', labelnames=['quantile'])
-
     def test_empty_labels_list(self):
         Histogram('h', 'help', [], registry=self.registry)
         self.assertEqual(0, self.registry.get_sample_value('h_sum'))
 
+    def test_unit_appended(self):
+        Histogram('h', 'help', [], registry=self.registry, unit="seconds")
+        self.assertEqual(0, self.registry.get_sample_value('h_seconds_sum'))
+
+    def test_unit_notappended(self):
+        Histogram('h_seconds', 'help', [], registry=self.registry, unit="seconds")
+        self.assertEqual(0, self.registry.get_sample_value('h_seconds_sum'))
+
+    def test_no_units_for_info_enum(self):
+        self.assertRaises(ValueError, Info, 'foo', 'help', unit="x")
+        self.assertRaises(ValueError, Enum, 'foo', 'help', unit="x")
+
     def test_wrapped_original_class(self):
         self.assertEqual(Counter.__wrapped__, Counter('foo', 'bar').__class__)
 
@@ -495,6 +506,10 @@ def test_untyped_labels(self):
         self.custom_collector(cmf)
         self.assertEqual(2, self.registry.get_sample_value('u', {'a': 'b', 'c': 'd'}))
 
+    def test_untyped_unit(self):
+        self.custom_collector(UntypedMetricFamily('u', 'help', value=1, unit='unit'))
+        self.assertEqual(1, self.registry.get_sample_value('u_unit', {}))
+
     def test_counter(self):
         self.custom_collector(CounterMetricFamily('c_total', 'help', value=1))
         self.assertEqual(1, self.registry.get_sample_value('c_total', {}))