coderanger
diff --git a/‎prometheus_client/core.py
Lines changed: 15 additions & 10 deletions b/‎prometheus_client/core.py
Lines changed: 15 additions & 10 deletions
diff --git a/‎prometheus_client/openmetrics/exposition.py
Lines changed: 4 additions & 2 deletions b/‎prometheus_client/openmetrics/exposition.py
Lines changed: 4 additions & 2 deletions
diff --git a/‎prometheus_client/openmetrics/parser.py
Lines changed: 27 additions & 4 deletions b/‎prometheus_client/openmetrics/parser.py
Lines changed: 27 additions & 4 deletions
diff --git a/‎tests/openmetrics/test_parser.py
Lines changed: 79 additions & 1 deletion b/‎tests/openmetrics/test_parser.py
Lines changed: 79 additions & 1 deletion
@@ -36,6 +36,7 @@
 _unpack_double = struct.Struct(b'd').unpack_from
 
 Sample = namedtuple('Sample', ['name', 'labels', 'value', 'timestamp', 'exemplar'])
+# Value can be an int or a float.
 # Timestamp and exemplar are optional.
 Sample.__new__.__defaults__ = (None, None)
 
@@ -168,8 +169,12 @@ class Metric(object):
     and SummaryMetricFamily instead.
     '''
     def __init__(self, name, documentation, typ, 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 not in _METRIC_TYPES:
             raise ValueError('Invalid metric type: ' + typ)
@@ -190,8 +195,8 @@ def __eq__(self, other):
                 self.samples == other.samples)
 
     def __repr__(self):
-        return "Metric(%s, %s, %s, %s)" % (self.name, self.documentation,
-            self.type, self.samples)
+        return "Metric(%s, %s, %s, %s, %s)" % (self.name, self.documentation,
+            self.type, self.unit, self.samples)
 
 class UntypedMetricFamily(Metric):
     '''A single untyped metric and its samples.
@@ -252,8 +257,8 @@ class GaugeMetricFamily(Metric):
 
     For use by custom collectors.
     '''
-    def __init__(self, name, documentation, value=None, labels=None):
-        Metric.__init__(self, name, documentation, 'gauge')
+    def __init__(self, name, documentation, value=None, labels=None, unit=''):
+        Metric.__init__(self, name, documentation, 'gauge', 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:
@@ -277,8 +282,8 @@ class SummaryMetricFamily(Metric):
 
     For use by custom collectors.
     '''
-    def __init__(self, name, documentation, count_value=None, sum_value=None, labels=None):
-        Metric.__init__(self, name, documentation, 'summary')
+    def __init__(self, name, documentation, count_value=None, sum_value=None, labels=None, unit=''):
+        Metric.__init__(self, name, documentation, 'summary', unit)
         if (sum_value is None) != (count_value is None):
             raise ValueError('count_value and sum_value must be provided together.')
         if labels is not None and count_value is not None:
@@ -306,8 +311,8 @@ class HistogramMetricFamily(Metric):
 
     For use by custom collectors.
     '''
-    def __init__(self, name, documentation, buckets=None, sum_value=None, labels=None):
-        Metric.__init__(self, name, documentation, 'histogram')
+    def __init__(self, name, documentation, buckets=None, sum_value=None, labels=None, unit=''):
+        Metric.__init__(self, name, documentation, 'histogram', unit)
         if (sum_value is None) != (buckets is None):
             raise ValueError('buckets and sum_value must be provided together.')
         if labels is not None and buckets is not None:
@@ -339,8 +344,8 @@ class GaugeHistogramMetricFamily(Metric):
 
     For use by custom collectors.
     '''
-    def __init__(self, name, documentation, buckets=None, labels=None):
-        Metric.__init__(self, name, documentation, 'gaugehistogram')
+    def __init__(self, name, documentation, buckets=None, labels=None, unit=''):
+        Metric.__init__(self, name, documentation, 'gaugehistogram', unit)
         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:
 
@@ -12,9 +12,11 @@ def generate_latest(registry):
     output = []
     for metric in registry.collect():
         mname = metric.name
-        output.append('# HELP {0} {1}'.format(
+        output.append('# HELP {0} {1}\n'.format(
             mname, metric.documentation.replace('\\', r'\\').replace('\n', r'\n').replace('"', r'\"')))
-        output.append('\n# TYPE {0} {1}\n'.format(mname, metric.type))
+        output.append('# TYPE {0} {1}\n'.format(mname, metric.type))
+        if metric.unit:
+            output.append('# UNIT {0} {1}\n'.format(mname, metric.unit))
         for s in metric.samples:
             if s.labels:
                 labelstr = '{{{0}}}'.format(','.join(
 
@@ -84,6 +84,8 @@ def _parse_sample(text):
             if char == '\\':
                 state = 'labelvalueslash'
             elif char == '"':
+                if not core._METRIC_LABEL_NAME_RE.match(''.join(labelname)):
+                    raise ValueError("Invalid line: " + text)
                 labels[''.join(labelname)] = ''.join(labelvalue)
                 labelname = []
                 labelvalue = []
@@ -120,8 +122,13 @@ def _parse_sample(text):
                 value.append(char)
     if not value:
         raise ValueError("Invalid line: " + text)
+    val = None
+    try:
+        val = int(''.join(value))
+    except ValueError:
+        val = float(''.join(value))
 
-    return core.Sample(''.join(name), labels, float(''.join(value)))
+    return core.Sample(''.join(name), labels, val)
 
 
 def text_fd_to_metric_families(fd):
@@ -141,11 +148,17 @@ def text_fd_to_metric_families(fd):
     allowed_names = []
     eof = False
 
+    seen_metrics = set()
     def build_metric(name, documentation, typ, unit, samples):
+        if name in seen_metrics:
+            raise ValueError("Duplicate metric: " + name)
+        seen_metrics.add(name)
         metric = core.Metric(name, documentation, typ, unit)
-        # TODO: chheck only hitogram buckets have exemplars.
+        # 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 metrics appear only once
+        # TODO: check for metadata in middle of samples
         metric.samples = samples
         return metric
 
@@ -160,7 +173,7 @@ def build_metric(name, documentation, typ, unit, samples):
             eof = True
         elif line.startswith('#'):
             parts = line.split(' ', 3)
-            if len(parts) < 2:
+            if len(parts) < 4:
                 raise ValueError("Invalid line: " + line)
             if parts[1] == 'HELP':
                 if parts[2] != name:
@@ -193,6 +206,16 @@ def build_metric(name, documentation, typ, unit, samples):
                     'gaugehistogram': ['_bucket'],
                 }.get(typ, [''])
                 allowed_names = [name + n for n in allowed_names]
+            elif parts[1] == 'UNIT':
+                if parts[2] != name:
+                    if name != '':
+                        yield build_metric(name, documentation, typ, unit, samples)
+                    # New metric
+                    name = parts[2]
+                    typ = 'untyped'
+                    samples = []
+                    allowed_names = [parts[2
F438
]]
+                unit = parts[3]
             else:
                 raise ValueError("Invalid line: " + line)
         else:
 
@@ -36,6 +36,14 @@ def test_simple_counter(self):
 """)
         self.assertEqual([CounterMetricFamily("a", "help", value=1)], list(families))
 
+    def test_uint64_counter(self):
+        families = text_string_to_metric_families("""# TYPE a counter
+# HELP a help
+a_total 9223372036854775808
+# EOF
+""")
+        self.assertEqual([CounterMetricFamily("a", "help", value=9223372036854775808)], list(families))
+
     def test_simple_gauge(self):
         families = text_string_to_metric_families("""# TYPE a gauge
 # HELP a help
@@ -44,6 +52,23 @@ def test_simple_gauge(self):
 """)
         self.assertEqual([GaugeMetricFamily("a", "help", value=1)], list(families))
 
+    def test_float_gauge(self):
+        families = text_string_to_metric_families("""# TYPE a gauge
+# HELP a help
+a 1.2
+# EOF
+""")
+        self.assertEqual([GaugeMetricFamily("a", "help", value=1.2)], list(families))
+
+    def test_unit_gauge(self):
+        families = text_string_to_metric_families("""# TYPE a_seconds gauge
+# UNIT a_seconds seconds
+# HELP a_seconds help
+a_seconds 1
+# EOF
+""")
+        self.assertEqual([GaugeMetricFamily("a_seconds", "help", value=1)], list(families))
+
     def test_simple_summary(self):
         families = text_string_to_metric_families("""# TYPE a summary
 # HELP a help
@@ -152,9 +177,15 @@ def test_nan(self):
         families = text_string_to_metric_families("""a NaN
 # EOF
 """)
-        # Can't use a simple comparison as nan != nan.
         self.assertTrue(math.isnan(list(families)[0].samples[0][2]))
 
+    def test_no_newline_after_eof(self):
+        families = text_string_to_metric_families("""# TYPE a gauge
+# HELP a help
+a 1
+# EOF""")
+        self.assertEqual([GaugeMetricFamily("a", "help", value=1)], list(families))
+
     def test_empty_label(self):
         families = text_string_to_metric_families("""# TYPE a counter
 # HELP a help
@@ -289,6 +320,53 @@ def collect(self):
         registry.register(TextCollector())
         self.assertEqual(text.encode('utf-8'), generate_latest(registry))
 
+    def test_invalid_input(self):
+        for case in [
+                # No EOF.
+                (''),
+                # Text after EOF.
+                ('a 1\n# EOF\nblah'),
+                ('a 1\n# EOFblah'),
+                # Missing or wrong quotes on label value.
+                ('a{a=1} 1\n# EOF\n'),
+                ('a{a="1} 1\n# EOF\n'),
+                ('a{a=\'1\'} 1\n# EOF\n'),
+                # Missing or extra commas.
+                ('a{a="1"b="2"} 1\n# EOF\n'),
+                ('a{a="1",,b="2"} 1\n# EOF\n'),
+                ('a{a="1",b="2",} 1\n# EOF\n'),
+                # Missing value.
+                ('a\n# EOF\n'),
+                ('a \n# EOF\n'),
+                # Bad HELP.
+                ('# HELP\n# EOF\n'),
+                ('# HELP \n# EOF\n'),
+                ('# HELP a\n# EOF\n'),
+                ('# HELP a\t\n# EOF\n'),
+                (' # HELP a meh\n# EOF\n'),
+                # Bad TYPE.
+                ('# TYPE\n# EOF\n'),
+                ('# TYPE \n# EOF\n'),
+                ('# TYPE a\n# EOF\n'),
+                ('# TYPE a\t\n# EOF\n'),
+                ('# TYPE a meh\n# EOF\n'),
+                ('# TYPE a meh \n# EOF\n'),
+                ('# TYPE a gauge \n# EOF\n'),
+                # Bad UNIT.
+                ('# UNIT\n# EOF\n'),
+                ('# UNIT \n# EOF\n'),
+                ('# UNIT a\n# EOF\n'),
+                ('# UNIT a\t\n# EOF\n'),
+                ('# UNIT a seconds\n# EOF\n'),
+                ('# UNIT a_seconds seconds \n# EOF\n'),
+                # Bad metric names.
+                ('0a 1\n# EOF\n'),
+                ('a.b 1\n# EOF\n'),
+                ('a-b 1\n# EOF\n'),
+                ]:
+            with self.assertRaises(ValueError):
+                list(text_string_to_metric_families(case))
+
 
 if __name__ == '__main__':
     unittest.main()