Add type annotations to synapse.metrics

synapse/metrics/__init__.py

@@ -24,11 +24,13 @@
     Any,
     Callable,
     Dict,
+    Generic,
     Iterable,
     Optional,
     Sequence,
     Set,
     Tuple,
+    Type,
     TypeVar,
     Union,
     cast,
@@ -113,11 +115,12 @@ def _register(self) -> None:
         all_gauges[self.name] = self
 
 
-# Placeholder for `InFlightGauge._metrics_class`, which is created dynamically
-_MetricsEntry = Any
+# `MetricsEntry` only makes sense when it is a `Protocol`,
+# but `Protocol` can't be used as a `TypeVar` bound.
+MetricsEntry = TypeVar("MetricsEntry")
 
 
-class InFlightGauge:
+class InFlightGauge(Generic[MetricsEntry]):
     """Tracks number of things (e.g. requests, Measure blocks, etc) in flight
     at any given time.
 
@@ -147,13 +150,13 @@ def __init__(
 
         # Create a class which have the sub_metrics values as attributes, which
         # default to 0 on initialization. Used to pass to registered callbacks.
-        self._metrics_class = attr.make_class(
+        self._metrics_class: Type[MetricsEntry] = attr.make_class(
             "_MetricsEntry", attrs={x: attr.ib(0) for x in sub_metrics}, slots=True
         )
 
         # Counts number of in flight blocks for a given set of label values
         self._registrations: Dict[
-            Tuple[str, ...], Set[Callable[[_MetricsEntry], None]]
+            Tuple[str, ...], Set[Callable[[MetricsEntry], None]]
         ] = {}
 
         # Protects access to _registrations
@@ -164,7 +167,7 @@ def __init__(
     def register(
         self,
         key: Tuple[str, ...],
-        callback: Callable[[_MetricsEntry], None],
+        callback: Callable[[MetricsEntry], None],
     ) -> None:
         """Registers that we've entered a new block with labels `key`.
 
@@ -184,7 +187,7 @@ def register(
     def unregister(
         self,
         key: Tuple[str, ...],
-        callback: Callable[[_MetricsEntry], None],
+        callback: Callable[[MetricsEntry], None],
     ) -> None:
         """Registers that we've exited a block with labels `key`."""
 

synapse/metrics/_exposition.py

@@ -25,7 +25,7 @@
 import threading
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from socketserver import ThreadingMixIn
-from typing import Any, Dict, List, Type
+from typing import Any, Dict, List, Type, Union
 from urllib.parse import parse_qs, urlparse
 
 from prometheus_client import REGISTRY, CollectorRegistry
@@ -39,15 +39,11 @@
 CONTENT_TYPE_LATEST = "text/plain; version=0.0.4; charset=utf-8"
 
 
-INF = float("inf")
-MINUS_INF = float("-inf")
-
-
-def floatToGoString(d: Any) -> str:
+def floatToGoString(d: Union[int, float]) -> str:
     d = float(d)
-    if d == INF:
+    if d == math.inf:
         return "+Inf"
-    elif d == MINUS_INF:
+    elif d == -math.inf:
         return "-Inf"
     elif math.isnan(d):
         return "NaN"

synapse/metrics/background_process_metrics.py

@@ -203,6 +203,9 @@ def run_as_background_process(  # type: ignore[misc]
     ...
     # The `type: ignore[misc]` above suppresses
     # "error: Overloaded function signatures 1 and 2 overlap with incompatible return types  [misc]"
+    #
+    # Overloads are used instead of a `Union` here because mypy fails to infer the type of `R`
+    # and complains about almost every call to this function when `func` is a `Union`.
 
 
 @overload
@@ -287,19 +290,19 @@ async def run() -> Optional[R]:
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-# NB: Return type is incorrect and should be a callable returning an F with a
-#     Deferred[Optional[R]] return, which we can't express correctly until Python 3.10.
 def wrap_as_background_process(desc: str) -> Callable[[F], F]:
     """Decorator that wraps a function that gets called as a background
     process.
 
     Equivalent to calling the function with `run_as_background_process`.
 
-    Note that `run_as_background_process` changes the return type into
+    Note that the annotated return type of this function is incorrect.
+    The return type of the function, once wrapped, is actually a
     `Deferred[Optional[T]]`.
     """
 
     # NB: Return type is incorrect and should be F with a Deferred[Optional[R]] return
+    #     We can't expression the return types correctly until Python 3.10.
     def wrap_as_background_process_inner(func: F) -> F:
         @wraps(func)
         def wrap_as_background_process_inner_2(

synapse/util/metrics.py

@@ -17,6 +17,7 @@
 from typing import Any, Callable, Optional, TypeVar, cast
 
 from prometheus_client import Counter
+from typing_extensions import Protocol
 
 from synapse.logging.context import (
     ContextResourceUsage,
@@ -53,8 +54,14 @@
     "synapse_util_metrics_block_db_sched_duration_seconds", "", ["block_name"]
 )
 
+
+class InFlightBlockMetrics(Protocol):
+    real_time_max: float
+    real_time_sum: float
+
+
 # Tracks the number of blocks currently active
-in_flight = InFlightGauge(
+in_flight = InFlightGauge[InFlightBlockMetrics](
     "synapse_util_metrics_block_in_flight",
     "",
     labels=["block_name"],
@@ -168,7 +175,7 @@ def get_resource_usage(self) -> ContextResourceUsage:
         """
         return self._logging_context.get_resource_usage()
 
-    def _update_in_flight(self, metrics):
+    def _update_in_flight(self, metrics: InFlightBlockMetrics):
         """Gets called when processing in flight metrics"""
         duration = self.clock.time() - self.start