Coerce label values to strings.

brian-brazil · brian-brazil · commit 72651128d04c · 2015-04-06T12:41:25.000+01:00
Make __all__ work with unicode_literals
diff --git a/prometheus_client/__init__.py b/prometheus_client/__init__.py
@@ -12,11 +12,14 @@
     from BaseHTTPServer import BaseHTTPRequestHandler
 except ImportError:
     # Python 3
+    unicode = str
     from http.server import BaseHTTPRequestHandler
 from functools import wraps
 from threading import Lock
 
 __all__ = ['Counter', 'Gauge', 'Summary', 'Histogram']
+# http://stackoverflow.com/questions/19913653/no-unicode-in-all-for-a-packages-init
+__all__ = [n.encode('ascii') for n in __all__]
 
 _METRIC_NAME_RE = re.compile(r'^[a-zA-Z_:][a-zA-Z0-9_:]*$')
 _METRIC_LABEL_NAME_RE = re.compile(r'^[a-zA-Z_:][a-zA-Z0-9_:]*$')
@@ -108,7 +111,7 @@ def labels(self, *labelvalues):
         '''Return the child for the given labelset.'''
         if len(labelvalues) != len(self._labelnames):
             raise ValueError('Incorrect label count')
-        labelvalues = tuple(labelvalues)
+        labelvalues = tuple([unicode(l) for l in labelvalues])
         with self._lock:
             if labelvalues not in self._metrics:
                 self._metrics[labelvalues] = self._wrappedClass(**self._kwargs)
@@ -118,7 +121,7 @@ def remove(self, *labelvalues):
         '''Remove the given labelset from the metric.'''
         if len(labelvalues) != len(self._labelnames):
             raise ValueError('Incorrect label count')
-        labelvalues = tuple(labelvalues)
+        labelvalues = tuple([unicode(l) for l in labelvalues])
         with self._lock:
             del self._metrics[labelvalues]
 
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -223,6 +223,18 @@ def test_incorrect_label_count_raises(self):
         self.assertRaises(ValueError, self.counter.remove)
         self.assertRaises(ValueError, self.counter.remove, 'a', 'b')
 
+    def test_labels_coerced_to_string(self):
+        self.counter.labels(None).inc()
+        self.assertEqual(1, self.registry.get_sample_value('c', {'l': 'None'}))
+
+        self.counter.remove(None)
+        self.assertEqual(None, self.registry.get_sample_value('c', {'l': 'None'}))
+
+    def test_non_string_labels_raises(self):
+        class Test(object):
+            __str__ = None
+        self.assertRaises(TypeError, self.counter.labels, Test())
+
     def test_namespace_subsystem_concatenated(self):
         c = Counter('c', 'help', namespace='a', subsystem='b', registry=self.registry)
         c.inc()
