Rework MetricConfig, make labels an attribute (#7)

README.rst

@@ -120,11 +120,14 @@ with a list of ``MetricConfig``\s. This is typically done in ``configure()``:
 
 .. code:: python
 
-    def configure(self, args: argparse.Namespace):
+    def configure(self, args: argparse.Namespace) -> None:
         # ...
         self.create_metrics(
-            [MetricConfig("metric1", "a metric", "gauge", {}),
-             MetricConfig("metric2", "another metric", "counter", {})])
+            [
+                MetricConfig("metric1", "a metric", "gauge"),
+                MetricConfig("metric2", "another metric", "counter", labels=("l1", "l2")),
+            ]
+        )
 
 
 Web application setup
@@ -145,11 +148,11 @@ coroutine and is called with a dict mapping metric names to metric objects:
 
 .. code:: python
 
-    async def on_application_startup(self, application):
+    async def on_application_startup(self, application: aiohttp.web.Application) -> None:
         # ...
         application["exporter"].set_metric_update_handler(self._update_handler)
 
-    async def _update_handler(self, metrics):
+    async def _update_handler(self, metrics: dict[str, prometheus_client.metrics.MetricWrapperBase]):
         for name, metric in metrics.items():
             metric.set(...)
 

prometheus_aioexporter/__init__.py

@@ -1,16 +1,18 @@
 """Asyncio library for creating Prometheus exporters."""
 
-from .metric import (
+from ._metric import (
+    InvalidMetricType,
     MetricConfig,
     MetricsRegistry,
 )
-from .script import PrometheusExporterScript
+from ._script import PrometheusExporterScript
 
 __all__ = [
-    "__version__",
+    "InvalidMetricType",
     "MetricConfig",
     "MetricsRegistry",
     "PrometheusExporterScript",
+    "__version__",
 ]
 
-__version__ = "1.7.0"
+__version__ = "2.0.dev1"

prometheus_aioexporter/metric.py → prometheus_aioexporter/_metric.py

@@ -1,14 +1,11 @@
-"""Helpers around prometheus-client to create and register metrics."""
+"""Helpers around prometheus_client to create and register metrics."""
 
 from collections.abc import Iterable
 from dataclasses import (
     dataclass,
     field,
 )
-from typing import (
-    Any,
-    NamedTuple,
-)
+from typing import Any
 
 from prometheus_client import (
     CollectorRegistry,
@@ -18,31 +15,28 @@
     generate_latest,
     Histogram,
     Info,
-    Metric,
     Summary,
 )
+from prometheus_client.metrics import MetricWrapperBase
 from prometheus_client.registry import Collector
 
 
-class MetricType(NamedTuple):
+@dataclass(frozen=True)
+class MetricType:
     """Details about a metric type."""
 
-    cls: Metric
-    options: dict[str, str] = field(default_factory=dict)
+    cls: type[MetricWrapperBase]
+    options: list[str] = field(default_factory=list)
 
 
 # Map metric types to their MetricTypes
 METRIC_TYPES: dict[str, MetricType] = {
-    "counter": MetricType(cls=Counter, options={"labels": "labelnames"}),
-    "enum": MetricType(
-        cls=Enum, options={"labels": "labelnames", "states": "states"}
-    ),
-    "gauge": MetricType(cls=Gauge, options={"labels": "labelnames"}),
-    "histogram": MetricType(
-        cls=Histogram, options={"labels": "labelnames", "buckets": "buckets"}
-    ),
-    "info": MetricType(cls=Info, options={"labels": "labelnames"}),
-    "summary": MetricType(cls=Summary, options={"labels": "labelnames"}),
+    "counter": MetricType(cls=Counter),
+    "enum": MetricType(cls=Enum, options=["states"]),
+    "gauge": MetricType(cls=Gauge),
+    "histogram": MetricType(cls=Histogram, options=["buckets"]),
+    "info": MetricType(cls=Info),
+    "summary": MetricType(cls=Summary),
 }
 
 
@@ -53,9 +47,11 @@ class MetricConfig:
     name: str
     description: str
     type: str
-    config: dict[str, Any]
+    labels: Iterable[str] = field(default_factory=tuple)
+    config: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
+        self.labels = tuple(sorted(self.labels))
         if self.type not in METRIC_TYPES:
             raise InvalidMetricType(self.name, self.type)
 
@@ -79,29 +75,29 @@ class MetricsRegistry:
 
     def __init__(self) -> None:
         self.registry = CollectorRegistry(auto_describe=True)
-        self._metrics: dict[str, Metric] = {}
+        self._metrics: dict[str, MetricWrapperBase] = {}
 
     def create_metrics(
         self, configs: Iterable[MetricConfig]
-    ) -> dict[str, Metric]:
+    ) -> dict[str, MetricWrapperBase]:
         """Create Prometheus metrics from a list of MetricConfigs."""
-        metrics: dict[str, Metric] = {
+        metrics: dict[str, MetricWrapperBase] = {
             config.name: self._register_metric(config) for config in configs
         }
         self._metrics.update(metrics)
         return metrics
 
     def get_metric(
         self, name: str, labels: dict[str, str] | None = None
-    ) -> Metric:
+    ) -> MetricWrapperBase:
         """Return a metric, optionally configured with labels."""
         metric = self._metrics[name]
         if labels:
             return metric.labels(**labels)
 
         return metric
 
-    def get_metrics(self) -> dict[str, Metric]:
+    def get_metrics(self) -> dict[str, MetricWrapperBase]:
         """Return a dict mapping names to metrics."""
         return self._metrics.copy()
 
@@ -118,13 +114,17 @@ def generate_metrics(self) -> bytes:
         """Generate text with metrics values from the registry."""
         return bytes(generate_latest(self.registry))
 
-    def _register_metric(self, config: MetricConfig) -> Metric:
+    def _register_metric(self, config: MetricConfig) -> MetricWrapperBase:
         metric_type = METRIC_TYPES[config.type]
         options = {
-            metric_type.options[key]: value
+            key: value
             for key, value in config.config.items()
             if key in metric_type.options
         }
         return metric_type.cls(
-            config.name, config.description, registry=self.registry, **options
+            config.name,
+            config.description,
+            labelnames=config.labels,
+            registry=self.registry,
+            **options,
         )

prometheus_aioexporter/script.py → prometheus_aioexporter/_script.py

@@ -8,18 +8,16 @@
 from typing import IO
 
 from aiohttp.web import Application
-from prometheus_client import (
-    Metric,
-    ProcessCollector,
-)
+from prometheus_client import ProcessCollector
+from prometheus_client.metrics import MetricWrapperBase
 from toolrack.log import setup_logger
 from toolrack.script import Script
 
-from .metric import (
+from ._metric import (
     MetricConfig,
     MetricsRegistry,
 )
-from .web import PrometheusExporter
+from ._web import PrometheusExporter
 
 
 class PrometheusExporterScript(Script):  # type: ignore
@@ -95,7 +93,7 @@ def configure(self, args: argparse.Namespace) -> None:
 
     def create_metrics(
         self, metric_configs: Iterable[MetricConfig]
-    ) -> dict[str, Metric]:
+    ) -> dict[str, MetricWrapperBase]:
         """Create and register metrics from a list of MetricConfigs."""
         return self.registry.create_metrics(metric_configs)
 
@@ -176,8 +174,10 @@ def _setup_logging(self, log_level: str) -> None:
     def _configure_registry(self, include_process_stats: bool = False) -> None:
         """Configure the MetricRegistry."""
         if include_process_stats:
+            # XXX ignore type until
+            # https://github.com/prometheus/client_python/pull/970 is fixed
             self.registry.register_additional_collector(
-                ProcessCollector(registry=None)
+                ProcessCollector(registry=None)  # type: ignore[arg-type]
             )
 
     def _get_ssl_context(

prometheus_aioexporter/web.py → prometheus_aioexporter/_web.py

@@ -3,7 +3,6 @@
 from collections.abc import (
     Awaitable,
     Callable,
-    Iterable,
 )
 from ssl import SSLContext
 from textwrap import dedent
@@ -14,15 +13,13 @@
     Response,
     run_app,
 )
-from prometheus_client import (
-    CONTENT_TYPE_LATEST,
-    Metric,
-)
+from prometheus_client import CONTENT_TYPE_LATEST
+from prometheus_client.metrics import MetricWrapperBase
 
-from .metric import MetricsRegistry
+from ._metric import MetricsRegistry
 
 # Signature for update handler
-UpdateHandler = Callable[[Iterable[Metric]], Awaitable[None]]
+UpdateHandler = Callable[[dict[str, MetricWrapperBase]], Awaitable[None]]
 
 
 class PrometheusExporter:
@@ -65,7 +62,7 @@ def set_metric_update_handler(self, handler: UpdateHandler) -> None:
         as argument, mapping metric names to metrics.  The signature is the
         following:
 
-          async def update_handler(metrics):
+          async def update_handler(metrics: dict[str, MetricWrapperBase]) -> None:
 
         """
         self._update_handler = handler

prometheus_aioexporter/sample.py

@@ -1,8 +1,13 @@
 from argparse import Namespace
 import random
+from typing import cast
 
 from aiohttp.web import Application
-from prometheus_client import Metric
+from prometheus_client import (
+    Counter,
+    Gauge,
+)
+from prometheus_client.metrics import MetricWrapperBase
 
 from . import (
     MetricConfig,
@@ -20,23 +25,27 @@ def configure(self, args: Namespace) -> None:
         self.create_metrics(
             [
                 MetricConfig(
-                    "a_gauge", "a gauge", "gauge", {"labels": ["foo", "bar"]}
+                    "a_gauge", "a gauge", "gauge", labels=("foo", "bar")
                 ),
                 MetricConfig(
-                    "a_counter", "a counter", "counter", {"labels": ["baz"]}
+                    "a_counter", "a counter", "counter", labels=("baz",)
                 ),
             ]
         )
 
     async def on_application_startup(self, application: Application) -> None:
         application["exporter"].set_metric_update_handler(self._update_handler)
 
-    async def _update_handler(self, metrics: dict[str, Metric]) -> None:
-        metrics["a_gauge"].labels(
+    async def _update_handler(
+        self, metrics: dict[str, MetricWrapperBase]
+    ) -> None:
+        gauge = cast(Gauge, metrics["a_gauge"])
+        gauge.labels(
             foo=random.choice(["this-foo", "other-foo"]),
             bar=random.choice(["this-bar", "other-bar"]),
         ).set(random.uniform(0, 100))
-        metrics["a_counter"].labels(
+        counter = cast(Counter, metrics["a_counter"])
+        counter.labels(
             baz=random.choice(["this-baz", "other-baz"]),
         ).inc(random.choice(range(10)))