Repository: prometheus/client_python (python)

Issue:
Registering a metric after creating RestrictedRegistry with its name does not add the metric to RestrictedRegistry
The method `restricted_registry` works by filtering all metrics once at the time of its creation, as one can see from looking at the file `registry.py`.

If you register a metric after creating the restriction for it, it is not added to the restricted registry - while the docstring could be interpreted as claiming otherwise.

Solution 1: the class could be documented in a way that makes its behaviour explicit;
Solution 2: the class could be expanded to recollect metrics from its parent on every `collect()` call.

Pull request:
[#674] Made the registry restriction handle newly added metrics.
Fixes #674.

Patch:
diff --git a/prometheus_client/metrics_core.py b/prometheus_client/metrics_core.py
index 4fd0fcd..4ab2af1 100644
--- a/prometheus_client/metrics_core.py
+++ b/prometheus_client/metrics_core.py
@@ -58,6 +58,15 @@ class Metric(object):
             self.samples,
         )
 
+    def _restricted_metric(self, names):
+        """Build a snapshot of a metric with samples restricted to a given set of names."""
+        samples = [s for s in self.samples if s[0] in names]
+        if samples:
+            m = Metric(self.name, self.documentation, self.type)
+            m.samples = samples
+            return m
+        return None
+
 
 class UnknownMetricFamily(Metric):
     """A single unknown metric and its samples.
diff --git a/prometheus_client/registry.py b/prometheus_client/registry.py
index fff1f98..0eb003b 100644
--- a/prometheus_client/registry.py
+++ b/prometheus_client/registry.py
@@ -94,28 +94,7 @@ class CollectorRegistry(object):
 
         Experimental."""
         names = set(names)
-        collectors = set()
-        metrics = []
-        with self._lock:
-            if 'target_info' in names and self._target_info:
-                metrics.append(self._target_info_metric())
-                names.remove('target_info')
-            for name in names:
-                if name in self._names_to_collectors:
-                    collectors.add(self._names_to_collectors[name])
-        for collector in collectors:
-            for metric in collector.collect():
-                samples = [s for s in metric.samples if s[0] in names]
-                if samples:
-                    m = Metric(metric.name, metric.documentation, metric.type)
-                    m.samples = samples
-                    metrics.append(m)
-
-        class RestrictedRegistry(object):
-            def collect(self):
-                return metrics
-
-        return RestrictedRegistry()
+        return RestrictedRegistry(names, self)
 
     def set_target_info(self, labels):
         with self._lock:
@@ -150,4 +129,16 @@ class CollectorRegistry(object):
         return None
 
 
+class RestrictedRegistry(object):
+    def __init__(self, names, registry):
+        self._name_set = set(names)
+        self._registry = registry
+
+    def collect(self):
+        for metric in self._registry.collect():
+            m = metric._restricted_metric(self._name_set)
+            if m:
+                yield m
+
+
 REGISTRY = CollectorRegistry(auto_describe=True)


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index dbb2703..d36aef0 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -796,7 +796,19 @@ class TestCollectorRegistry(unittest.TestCase):
 
         m = Metric('s', 'help', 'summary')
         m.samples = [Sample('s_sum', {}, 7)]
-        self.assertEqual([m], registry.restricted_registry(['s_sum']).collect())
+        self.assertEqual([m], list(registry.restricted_registry(['s_sum']).collect()))
+
+    def test_restricted_registry_adds_new_metrics(self):
+        registry = CollectorRegistry()
+        Counter('c_total', 'help', registry=registry)
+
+        restricted_registry = registry.restricted_registry(['s_sum'])
+
+        Summary('s', 'help', registry=registry).observe(7)
+        m = Metric('s', 'help', 'summary')
+        m.samples = [Sample('s_sum', {}, 7)]
+
+        self.assertEqual([m], list(restricted_registry.collect()))
 
     def test_target_info_injected(self):
         registry = CollectorRegistry(target_info={'foo': 'bar'})
@@ -820,11 +832,11 @@ class TestCollectorRegistry(unittest.TestCase):
 
         m = Metric('s', 'help', 'summary')
         m.samples = [Sample('s_sum', {}, 7)]
-        self.assertEqual([m], registry.restricted_registry(['s_sum']).collect())
+        self.assertEqual([m], list(registry.restricted_registry(['s_sum']).collect()))
 
         m = Metric('target', 'Target metadata', 'info')
         m.samples = [Sample('target_info', {'foo': 'bar'}, 1)]
-        self.assertEqual([m], registry.restricted_registry(['target_info']).collect())
+        self.assertEqual([m], list(registry.restricted_registry(['target_info']).collect()))
 
 
 if __name__ == '__main__':