Allow to add labels inside a context manager

prometheus_client/context_managers.py

@@ -42,19 +42,26 @@ def wrapped(func, *args, **kwargs):
 
 
 class Timer:
-    def __init__(self, callback):
-        self._callback = callback
+    def __init__(self, metric, callback_name):
+        self._metric = metric
+        self._callback_name = callback_name
 
     def _new_timer(self):
-        return self.__class__(self._callback)
+        return self.__class__(self._metric, self._callback_name)
 
     def __enter__(self):
         self._start = default_timer()
+        return self
 
     def __exit__(self, typ, value, traceback):
         # Time can go backwards.
         duration = max(default_timer() - self._start, 0)
-        self._callback(duration)
+        self._metric._raise_if_not_observable()
+        callback = getattr(self._metric, self._callback_name)
+        callback(duration)
+
+    def labels(self, *args, **kw):
+        self._metric = self._metric.labels(*args, **kw)
 
     def __call__(self, f):
         def wrapped(func, *args, **kwargs):

prometheus_client/metrics.py

@@ -396,8 +396,7 @@ def time(self):
 
         Can be used as a function decorator or context manager.
         """
-        self._raise_if_not_observable()
-        return Timer(self.set)
+        return Timer(self, 'set')
 
     def set_function(self, f):
         """Call the provided function to return the Gauge value.
@@ -475,8 +474,7 @@ def time(self):
 
         Can be used as a function decorator or context manager.
         """
-        self._raise_if_not_observable()
-        return Timer(self.observe)
+        return Timer(self, 'observe')
 
     def _child_samples(self):
         return (
@@ -600,7 +598,7 @@ def time(self):
 
         Can be used as a function decorator or context manager.
         """
-        return Timer(self.observe)
+        return Timer(self, 'observe')
 
     def _child_samples(self):
         samples = []

tests/test_core.py

@@ -222,13 +222,25 @@ def test_time_block_decorator(self):
             time.sleep(.001)
         self.assertNotEqual(0, self.registry.get_sample_value('g'))
 
+    def test_time_block_decorator_with_label(self):
+        value = self.registry.get_sample_value
+        self.assertEqual(None, value('g2', {'label1': 'foo'}))
+        with self.gauge_with_label.time() as metric:
+            metric.labels('foo')
+        self.assertLess(0, value('g2', {'label1': 'foo'}))
+
     def test_track_in_progress_not_observable(self):
         g = Gauge('test', 'help', labelnames=('label',), registry=self.registry)
         assert_not_observable(g.track_inprogress)
 
     def test_timer_not_observable(self):
         g = Gauge('test', 'help', labelnames=('label',), registry=self.registry)
-        assert_not_observable(g.time)
+
+        def manager():
+            with g.time():
+                pass
+
+        assert_not_observable(manager)
 
 
 class TestSummary(unittest.TestCase):
@@ -318,10 +330,21 @@ def test_block_decorator(self):
             pass
         self.assertEqual(1, self.registry.get_sample_value('s_count'))
 
+    def test_block_decorator_with_label(self):
+        value = self.registry.get_sample_value
+        self.assertEqual(None, value('s_with_labels_count', {'label1': 'foo'}))
+        with self.summary_with_labels.time() as metric:
+            metric.labels('foo')
+        self.assertEqual(1, value('s_with_labels_count', {'label1': 'foo'}))
+
     def test_timer_not_observable(self):
         s = Summary('test', 'help', labelnames=('label',), registry=self.registry)
 
-        assert_not_observable(s.time)
+        def manager():
+            with s.time():
+                pass
+
+        assert_not_observable(manager)
 
 
 class TestHistogram(unittest.TestCase):
@@ -435,6 +458,15 @@ def test_block_decorator(self):
         self.assertEqual(1, self.registry.get_sample_value('h_count'))
         self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
 
+    def test_block_decorator_with_label(self):
+        value = self.registry.get_sample_value
+        self.assertEqual(None, value('hl_count', {'l': 'a'}))
+        self.assertEqual(None, value('hl_bucket', {'le': '+Inf', 'l': 'a'}))
+        with self.labels.time() as metric:
+            metric.labels('a')
+        self.assertEqual(1, value('hl_count', {'l': 'a'}))
+        self.assertEqual(1, value('hl_bucket', {'le': '+Inf', 'l': 'a'}))
+
     def test_exemplar_invalid_label_name(self):
         self.assertRaises(ValueError, self.histogram.observe, 3.0, exemplar={':o)': 'smile'})
         self.assertRaises(ValueError, self.histogram.observe, 3.0, exemplar={'1': 'number'})