Annotate few functions and methods

Signed-off-by: Derek Kulinski <d@kulinski.us>

Author: takeda

prometheus_client/context_managers.py

@@ -1,21 +1,28 @@
 from timeit import default_timer
+from types import TracebackType
+from typing import Any, Callable, Optional, Type, TYPE_CHECKING, TypeVar
 
 from .decorator import decorate
 
+if TYPE_CHECKING:
+    from . import Counter
+    F = TypeVar("F", bound=Callable[..., Any])
+
 
 class ExceptionCounter:
-    def __init__(self, counter, exception):
+    def __init__(self, counter: "Counter", exception: Type[BaseException]) -> None:
         self._counter = counter
         self._exception = exception
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         pass
 
-    def __exit__(self, typ, value, traceback):
+    def __exit__(self, typ: Optional[Type[BaseException]], value: Optional[BaseException], traceback: Optional[TracebackType]) -> bool:
         if isinstance(value, self._exception):
             self._counter.inc()
+        return False
 
-    def __call__(self, f):
+    def __call__(self, f: "F") -> "F":
         def wrapped(func, *args, **kwargs):
             with self:
                 return func(*args, **kwargs)

prometheus_client/decorator.py

@@ -39,6 +39,9 @@
 import operator
 import re
 import sys
+from typing import Any, Callable, TypeVar
+
+F = TypeVar("F", bound=Callable[..., Any])
 
 __version__ = '4.0.10'
 
@@ -226,7 +229,7 @@ def create(cls, obj, body, evaldict, defaults=None,
                          evaldict, addsource, **attrs)
 
 
-def decorate(func, caller):
+def decorate(func: F, caller: F) -> F:
     """
     decorate(func, caller) decorates a function using a caller.
     """

prometheus_client/metrics.py

@@ -1,17 +1,23 @@
 from threading import Lock
 import time
 import types
+from typing import (
+    Any, Callable, Dict, Optional, Sequence, Tuple, Type, TypeVar,
+)
 
 from . import values  # retain this import style for testability
 from .context_managers import ExceptionCounter, InprogressTracker, Timer
 from .metrics_core import (
     Metric, METRIC_LABEL_NAME_RE, METRIC_NAME_RE,
     RESERVED_METRIC_LABEL_NAME_RE,
 )
-from .registry import REGISTRY
+from .registry import CollectorRegistry, REGISTRY
 from .samples import Exemplar
 from .utils import floatToGoString, INF
 
+T = TypeVar('T', bound='MetricWrapperBase')
+F = TypeVar("F", bound=Callable[..., Any])
+
 
 def _build_full_name(metric_type, name, namespace, subsystem, unit):
     full_name = ''
@@ -95,15 +101,15 @@ def __repr__(self):
         return f"{metric_type.__module__}.{metric_type.__name__}({self._name})"
 
     def __init__(self,
-                 name,
-                 documentation,
-                 labelnames=(),
-                 namespace='',
-                 subsystem='',
-                 unit='',
-                 registry=REGISTRY,
-                 _labelvalues=None,
-                 ):
+                 name: str,
+                 documentation: str,
+                 labelnames: Sequence[str]=(),
+                 namespace: str='',
+                 subsystem: str='',
+                 unit: str='',
+                 registry: CollectorRegistry=REGISTRY,
+                 _labelvalues: Optional[Sequence[str]]=None,
+                 ) -> None:
         self._name = _build_full_name(self._type, name, namespace, subsystem, unit)
         self._labelnames = _validate_labelnames(self, labelnames)
         self._labelvalues = tuple(_labelvalues or ())
@@ -127,7 +133,7 @@ def __init__(self,
             if registry:
                 registry.register(self)
 
-    def labels(self, *labelvalues, **labelkwargs):
+    def labels(self: T, *labelvalues: str, **labelkwargs: str) -> T:
         """Return the child for the given labelset.
 
         All metrics can have labels, allowing grouping of related time series.
@@ -193,7 +199,7 @@ def remove(self, *labelvalues):
         with self._lock:
             del self._metrics[labelvalues]
 
-    def clear(self):
+    def clear(self) -> None:
         """Remove all labelsets from the metric"""
         with self._lock:
             self._metrics = {}
@@ -213,6 +219,7 @@ def _multi_samples(self):
                 yield (suffix, dict(series_labels + list(sample_labels.items())), value, timestamp, exemplar)
 
     def _child_samples(self):  # pragma: no cover
+        # type: () -> Sequence[Tuple[str, Dict[str, str], float]]
         raise NotImplementedError('_child_samples() must be implemented by %r' % self)
 
     def _metric_init(self):  # pragma: no cover
@@ -258,12 +265,12 @@ def f():
     """
     _type = 'counter'
 
-    def _metric_init(self):
+    def _metric_init(self) -> None:
         self._value = values.ValueClass(self._type, self._name, self._name + '_total', self._labelnames,
                                         self._labelvalues)
         self._created = time.time()
 
-    def inc(self, amount=1, exemplar=None):
+    def inc(self, amount: float=1, exemplar: Optional[Dict[str, str]]=None) -> None:
         """Increment counter by the given amount."""
         self._raise_if_not_observable()
         if amount < 0:
@@ -273,7 +280,7 @@ def inc(self, amount=1, exemplar=None):
             _validate_exemplar(exemplar)
             self._value.set_exemplar(Exemplar(exemplar, amount, time.time()))
 
-    def count_exceptions(self, exception=Exception):
+    def count_exceptions(self, exception: Type[BaseException]=Exception) -> ExceptionCounter:
         """Count exceptions in a block of code or function.
 
         Can be used as a function decorator or context manager.
@@ -667,15 +674,15 @@ class Enum(MetricWrapperBase):
     _type = 'stateset'
 
     def __init__(self,
-                 name,
-                 documentation,
-                 labelnames=(),
-                 namespace='',
-                 subsystem='',
-                 unit='',
-                 registry=REGISTRY,
-                 _labelvalues=None,
-                 states=None,
+                 name: str,
+                 documentation: str,
+                 labelnames: Sequence[str]=(),
+                 namespace: str='',
+                 subsystem: str='',
+                 unit: str='',
+                 registry: CollectorRegistry=REGISTRY,
+                 _labelvalues: Optional[Sequence[str]]=None,
+                 states: Optional[Sequence[str]]=None,
                  ):
         super().__init__(
             name=name,
@@ -693,11 +700,11 @@ def __init__(self,
             raise ValueError(f'No states provided for Enum metric: {name}')
         self._kwargs['states'] = self._states = states
 
-    def _metric_init(self):
+    def _metric_init(self) -> None:
         self._value = 0
         self._lock = Lock()
 
-    def state(self, state):
+    def state(self, state: str) -> None:
         """Set enum metric state."""
         self._raise_if_not_observable()
         with self._lock:

prometheus_client/py.typed

@@ -23,6 +23,9 @@
         'prometheus_client.openmetrics',
         'prometheus_client.twisted',
     ],
+    package_data={
+        'prometheus_client': ['py.typed']
+    },
     extras_require={
         'twisted': ['twisted'],
     },