Preserve unit metadata in restricted registries

prometheus_client/metrics_core.py

@@ -59,7 +59,7 @@ 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 = Metric(self.name, self.documentation, self.type, self.unit)
             m.samples = samples
             return m
         return None

tests/test_core.py

@@ -1115,6 +1115,14 @@ def test_restricted_registry_adds_new_metrics(self):
 
         self.assertEqual([m], list(restricted_registry.collect()))
 
+    def test_restricted_registry_preserves_unit(self):
+        registry = CollectorRegistry()
+        Summary('s', 'help', unit='seconds', registry=registry).observe(7)
+
+        m = Metric('s_seconds', 'help', 'summary', unit='seconds')
+        m.samples = [Sample('s_seconds_sum', {}, 7)]
+        self.assertEqual([m], list(registry.restricted_registry(['s_seconds_sum']).collect()))
+
     def test_target_info_injected(self):
         registry = CollectorRegistry(target_info={'foo': 'bar'})
         self.assertEqual(1, registry.get_sample_value('target_info', {'foo': 'bar'}))

tests/test_wsgi.py

@@ -57,6 +57,26 @@ def validate_metrics(self, metric_name, help_text, increments):
     def test_report_metrics_1(self):
         self.validate_metrics("counter", "A counter", 2)
 
+    def test_filtered_openmetrics_preserves_unit(self):
+        counter = Counter('response_size', 'Response size', unit='bytes', registry=self.registry)
+        counter.inc(7)
+        self.environ['HTTP_ACCEPT'] = 'application/openmetrics-text; version=1.0.0'
+        self.environ['QUERY_STRING'] = 'name[]=response_size_bytes_total'
+
+        output = make_wsgi_app(self.registry)(self.environ, self.capture)
+
+        self.assertEqual(self.captured_status, '200 OK')
+        self.assertEqual(
+            output,
+            [
+                b'# HELP response_size_bytes Response size\n'
+                b'# TYPE response_size_bytes counter\n'
+                b'# UNIT response_size_bytes bytes\n'
+                b'response_size_bytes_total 7.0\n'
+                b'# EOF\n'
+            ],
+        )
+
     def test_report_metrics_2(self):
         self.validate_metrics("counter", "Another counter", 3)