From d985bb3c9b66446568b15aef9ae408e68ef56208 Mon Sep 17 00:00:00 2001
From: Andreas Zeidler <andreas.zeidler@zeit.de>
Date: Sat, 11 Dec 2021 13:40:20 +0100
Subject: [PATCH] Pass metric instance and callback name to `Timer`

This should make the code slightly more readable.
---
 prometheus_client/context_managers.py | 18 ++++++++----------
 prometheus_client/metrics.py          |  6 +++---
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/prometheus_client/context_managers.py b/prometheus_client/context_managers.py
index eba60bcb..c18a6e34 100644
--- a/prometheus_client/context_managers.py
+++ b/prometheus_client/context_managers.py
@@ -42,11 +42,12 @@ 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()
@@ -55,15 +56,12 @@ def __enter__(self):
     def __exit__(self, typ, value, traceback):
         # Time can go backwards.
         duration = max(default_timer() - self._start, 0)
-        instance = self._callback.__self__
-        instance._raise_if_not_observable()
-        self._callback(duration)
+        self._metric._raise_if_not_observable()
+        callback = getattr(self._metric, self._callback_name)
+        callback(duration)
 
     def labels(self, *args, **kw):
-        instance = self._callback.__self__
-        self._callback = getattr(
-            instance.labels(*args, **kw),
-            self._callback.__name__)
+        self._metric = self._metric.labels(*args, **kw)
 
     def __call__(self, f):
         def wrapped(func, *args, **kwargs):
diff --git a/prometheus_client/metrics.py b/prometheus_client/metrics.py
index 252b859e..802adeb7 100644
--- a/prometheus_client/metrics.py
+++ b/prometheus_client/metrics.py
@@ -396,7 +396,7 @@ def time(self):
 
         Can be used as a function decorator or context manager.
         """
-        return Timer(self.set)
+        return Timer(self, 'set')
 
     def set_function(self, f):
         """Call the provided function to return the Gauge value.
@@ -474,7 +474,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):
         return (
@@ -598,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 = []