Annotate few functions and methods

Author: Derek Kulinski

prometheus_client/context_managers.py

@@ -1,23 +1,34 @@
 from __future__ import unicode_literals
 
 from timeit import default_timer
+from types import TracebackType
+from typing import Any, Callable, Optional, TYPE_CHECKING, Type, TypeVar
 
 from .decorator import decorate
 
+if TYPE_CHECKING:
+    from . import Counter
+    F = TypeVar("F", bound=Callable[..., Any])
+
 
 class ExceptionCounter(object):
     def __init__(self, counter, exception):
+        # type: (Counter, Type[BaseException]) -> None
         self._counter = counter
         self._exception = exception
 
     def __enter__(self):
+        # type: () -> None
         pass
 
     def __exit__(self, typ, value, traceback):
+        # type: (Optional[Type[BaseException]], Optional[BaseException] , Optional[TracebackType]) -> bool
         if isinstance(value, self._exception):
             self._counter.inc()
+        return False
 
     def __call__(self, f):
+        # type: (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 TypeVar, Callable, Any
+
+F = TypeVar("F", bound=Callable[..., Any])
 
 __version__ = '4.0.10'
 
@@ -227,6 +230,7 @@ def create(cls, obj, body, evaldict, defaults=None,
 
 
 def decorate(func, caller):
+    # type: (F, F) -> F
     """
     decorate(func, caller) decorates a function using a caller.
     """

prometheus_client/metrics.py

@@ -2,16 +2,20 @@
 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 .utils import floatToGoString, INF
 
+T = TypeVar('T', bound='MetricWrapperBase')
+F = TypeVar("F", bound=Callable[..., Any])
+
 if sys.version_info > (3,):
     unicode = str
     create_bound_method = types.MethodType
@@ -97,6 +101,7 @@ def __init__(self,
                  registry=REGISTRY,
                  _labelvalues=None,
                  ):
+        # type: (str, str, Sequence[str], str, str, str, CollectorRegistry, Optional[Sequence[str]]) -> None
         self._name = _build_full_name(self._type, name, namespace, subsystem, unit)
         self._labelnames = _validate_labelnames(self, labelnames)
         self._labelvalues = tuple(_labelvalues or ())
@@ -121,6 +126,7 @@ def __init__(self,
                 registry.register(self)
 
     def labels(self, *labelvalues, **labelkwargs):
+        # type: (T, *str, **str) -> T
         """Return the child for the given labelset.
 
         All metrics can have labels, allowing grouping of related time series.
@@ -187,6 +193,7 @@ def remove(self, *labelvalues):
             del self._metrics[labelvalues]
 
     def clear(self):
+        # type: () -> None
         """Remove all labelsets from the metric"""
         with self._lock:
             self._metrics = {}
@@ -206,6 +213,7 @@ def _multi_samples(self):
                 yield (suffix, dict(series_labels + list(sample_labels.items())), value)
 
     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
@@ -252,18 +260,21 @@ def f():
     _type = 'counter'
 
     def _metric_init(self):
+        # type: () -> 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):
+        # type: (float) -> None
         """Increment counter by the given amount."""
         self._raise_if_not_observable()
         if amount < 0:
             raise ValueError('Counters can only be incremented by non-negative amounts.')
         self._value.inc(amount)
 
     def count_exceptions(self, exception=Exception):
+        # type: (Type[BaseException]) -> ExceptionCounter
         """Count exceptions in a block of code or function.
 
         Can be used as a function decorator or context manager.
@@ -663,6 +674,7 @@ def __init__(self,
                  _labelvalues=None,
                  states=None,
                  ):
+        # type: (str, str, Sequence[str], str, str, str, CollectorRegistry, Optional[Sequence[str]], Optional[Sequence[str]]) -> None
         super(Enum, self).__init__(
             name=name,
             documentation=documentation,
@@ -684,6 +696,7 @@ def _metric_init(self):
         self._lock = Lock()
 
     def state(self, state):
+        # type: (str) -> None
         """Set enum metric state."""
         self._raise_if_not_observable()
         with self._lock:

prometheus_client/py.typed

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