Expose original metric classes via __wrapped__ attribute (prometheus#151)

beezz · brian-brazil · commit c2e4bdae72d4 · 2017-03-24T17:39:10.000Z
Expose original metric classes via __wrapped__ attribute
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -556,6 +556,7 @@ def collect():
             registry.register(collector)
         return collector
 
+    init.__wrapped__ = cls
     return init
 
 
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -307,6 +307,9 @@ def test_empty_labels_list(self):
         h = Histogram('h', 'help', [], registry=self.registry)
         self.assertEqual(0, self.registry.get_sample_value('h_sum'))
 
+    def test_wrapped_original_class(self):
+        self.assertEqual(Counter.__wrapped__, Counter('foo', 'bar').__class__)
+
 
 class TestMetricFamilies(unittest.TestCase):
     def setUp(self):
@@ -386,6 +389,7 @@ def test_bad_constructors(self):
         self.assertRaises(ValueError, HistogramMetricFamily, 'h', 'help', buckets={}, sum_value=1, labels=['a'])
         self.assertRaises(KeyError, HistogramMetricFamily, 'h', 'help', buckets={}, sum_value=1)
 
+
 class TestCollectorRegistry(unittest.TestCase):
     def test_duplicate_metrics_raises(self):
         registry = CollectorRegistry()
