Annotate few functions and methods (#705)

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

prometheus_client/context_managers.py

@@ -1,21 +1,30 @@
 from timeit import default_timer
+from types import TracebackType
+from typing import (
+    Any, Callable, Literal, 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]) -> Literal[False]:
         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/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 = ''
@@ -56,8 +62,8 @@ def _validate_exemplar(exemplar):
 
 
 class MetricWrapperBase:
-    _type = None
-    _reserved_labelnames = ()
+    _type: Optional[str] = None
+    _reserved_labelnames: Sequence[str] = ()
 
     def _is_observable(self):
         # Whether this metric is observable, i.e.
@@ -94,20 +100,20 @@ def __repr__(self):
         metric_type = type(self)
         return f"{metric_type.__module__}.{metric_type.__name__}({self._name})"
 
-    def __init__(self,
-                 name,
-                 documentation,
-                 labelnames=(),
-                 namespace='',
-                 subsystem='',
-                 unit='',
-                 registry=REGISTRY,
-                 _labelvalues=None,
-                 ):
+    def __init__(self: T,
+                 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 ())
-        self._kwargs = {}
+        self._kwargs: Dict[str, Any] = {}
         self._documentation = documentation
         self._unit = unit
 
@@ -117,7 +123,7 @@ def __init__(self,
         if self._is_parent():
             # Prepare the fields needed for child metrics.
             self._lock = Lock()
-            self._metrics = {}
+            self._metrics: Dict[Sequence[str], T] = {}
 
         if self._is_observable():
             self._metric_init()
@@ -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 = {}
@@ -212,7 +218,7 @@ def _multi_samples(self):
             for suffix, sample_labels, value, timestamp, exemplar in metric._samples():
                 yield (suffix, dict(series_labels + list(sample_labels.items())), value, timestamp, exemplar)
 
-    def _child_samples(self):  # pragma: no cover
+    def _child_samples(self) -> Sequence[Tuple[str, Dict[str, str], float]]:  # pragma: no cover
         raise NotImplementedError('_child_samples() must be implemented by %r' % self)
 
     def _metric_init(self):  # pragma: no cover
@@ -258,12 +264,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 +279,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 +673,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 +699,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/openmetrics/parser.py

@@ -1,19 +1,14 @@
 #!/usr/bin/python
 
 
+import io as StringIO
 import math
 import re
 
 from ..metrics_core import Metric, METRIC_LABEL_NAME_RE
 from ..samples import Exemplar, Sample, Timestamp
 from ..utils import floatToGoString
 
-try:
-    import StringIO
-except ImportError:
-    # Python 3
-    import io as StringIO
-
 
 def text_string_to_metric_families(text):
     """Parse Openmetrics text format from a unicode string.

prometheus_client/parser.py

@@ -1,14 +1,9 @@
+import io as StringIO
 import re
 
 from .metrics_core import Metric
 from .samples import Sample
 
-try:
-    import StringIO
-except ImportError:
-    # Python 3
-    import io as StringIO
-
 
 def text_string_to_metric_families(text):
     """Parse Prometheus text format from a unicode string.

prometheus_client/samples.py

@@ -1,4 +1,4 @@
-from collections import namedtuple
+from typing import Dict, NamedTuple, Optional
 
 
 class Timestamp:
@@ -36,8 +36,15 @@ def __gt__(self, other):
 # Timestamp can be a float containing a unixtime in seconds,
 # a Timestamp object, or None.
 # Exemplar can be an Exemplar object, or None.
-Sample = namedtuple('Sample', ['name', 'labels', 'value', 'timestamp', 'exemplar'])
-Sample.__new__.__defaults__ = (None, None)
-
-Exemplar = namedtuple('Exemplar', ['labels', 'value', 'timestamp'])
-Exemplar.__new__.__defaults__ = (None,)
+class Exemplar(NamedTuple):
+    labels: Dict[str, str]
+    value: float
+    timestamp: Optional[float] = None
+
+
+class Sample(NamedTuple):
+    name: str
+    labels: Dict[str, str]
+    value: float
+    timestamp: Optional[float] = None
+    exemplar: Optional[Exemplar] = None

setup.py

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