Extend type forwarding to other context managers (prometheus#903)

crutcher · web-flow · commit c48e26e2283f · 2023-03-21T16:26:35.000-06:00
Signed-off-by: Crutcher Dunnavant &lt;crutcher@openai.com&gt;
diff --git a/prometheus_client/context_managers.py b/prometheus_client/context_managers.py
@@ -46,7 +46,7 @@ def __enter__(self):
     def __exit__(self, typ, value, traceback):
         self._gauge.dec()
 
-    def __call__(self, f):
+    def __call__(self, f: "F") -> "F":
         def wrapped(func, *args, **kwargs):
             with self:
                 return func(*args, **kwargs)
@@ -75,7 +75,7 @@ def __exit__(self, typ, value, traceback):
     def labels(self, *args, **kw):
         self._metric = self._metric.labels(*args, **kw)
 
-    def __call__(self, f):
+    def __call__(self, f: "F") -> "F":
         def wrapped(func, *args, **kwargs):
             # Obtaining new instance of timer every time
             # ensures thread safety and reentrancy.
