From 37f4f0383dc94ea9abd1b62ccfc44cdd900997e2 Mon Sep 17 00:00:00 2001 From: Ajay Kumar <206070437+ajaykr0905@users.noreply.github.com> Date: Mon, 19 Oct 2026 21:57:22 +0530 Subject: [PATCH] Support Python enum classes as metric states Accept enum classes as Enum metric states and export canonical member names. Preserve the class for labelled children and validate member selection, aliases and empty definitions. Refs #975 Assisted-by: OpenAI Codex Signed-off-by: Ajay Kumar <206070437+ajaykr0905@users.noreply.github.com> --- docs/content/instrumenting/enum.md | 28 +++++++++++++- prometheus_client/metrics.py | 22 +++++++++-- tests/test_core.py | 59 ++++++++++++++++++++++++++++++ 3 files changed, 103 insertions(+), 6 deletions(-) diff --git a/docs/content/instrumenting/enum.md b/docs/content/instrumenting/enum.md index b1e6169a..77803146 100644 --- a/docs/content/instrumenting/enum.md +++ b/docs/content/instrumenting/enum.md @@ -17,6 +17,27 @@ Enum exposes one time series per state: The first listed state is the default. +You can also pass a Python `enum.Enum` class. The member **names** become label +values; the underlying values may be strings, integers, or other objects. +Set the current state using a member of that class or its name: + +```python +import enum +from prometheus_client import Enum + +class TaskState(enum.Enum): + starting = 1 + running = 2 + stopped = 3 + +e = Enum('task_state', 'Task lifecycle', states=TaskState) +e.state(TaskState.running) # exports task_state{task_state="running"} 1.0 +e.state('stopped') +``` + +The first member is the default. Aliases refer to the same state and do not +create additional time series. An empty enum class raises `ValueError`. + Note: Enum metrics do not work in multiprocess mode. ## Constructor @@ -34,7 +55,7 @@ Enum(name, documentation, labelnames=(), namespace='', subsystem='', unit='', re | `subsystem` | `str` | `''` | Optional middle component. | | `unit` | `str` | `''` | Not supported — raises `ValueError`. Enum metrics cannot have a unit. | | `registry` | `CollectorRegistry` | `REGISTRY` | Registry to register with. Pass `None` to skip registration, which is useful in tests where you create metrics without wanting them in the global registry. | -| `states` | `List[str]` | required | The complete list of valid states. Must be non-empty. The first entry is the initial state. | +| `states` | `Sequence[str]` or `Type[enum.Enum]` | required | The complete set of valid states. Must be non-empty. The first entry or member is the initial state. | `namespace`, `subsystem`, and `name` are joined with underscores to form the full metric name: @@ -48,7 +69,10 @@ Enum('state', 'Worker state', states=['idle', 'running', 'error'], namespace='my ### `state(state)` -Set the current state. The value must be one of the strings passed in the `states` list. Raises `ValueError` if the state is not recognized. +Set the current state. The value must be one of the strings passed in the `states` +sequence. When `states` is a Python enum class, use a member of that class or its +name. A member of a different enum class is rejected, even if it has the same +name or value. Raises `ValueError` if the state is not recognized. ```python e.state('running') diff --git a/prometheus_client/metrics.py b/prometheus_client/metrics.py index dfa158b9..60b12911 100644 --- a/prometheus_client/metrics.py +++ b/prometheus_client/metrics.py @@ -1,3 +1,4 @@ +from enum import Enum as PythonEnum import os from threading import Lock import time @@ -761,6 +762,8 @@ class Enum(MetricWrapperBase): e.state('running') The first listed state will be the default. + A Python enum.Enum class can also supply the states. Its member names are + exported as labels, and state() accepts its members or their names. Enum metrics do not work in multiprocess mode. """ _type = 'stateset' @@ -774,11 +777,17 @@ def __init__(self, unit: str = '', registry: Optional[CollectorRegistry] = REGISTRY, _labelvalues: Optional[Sequence[str]] = None, - states: Optional[Sequence[str]] = None, + states: Optional[Union[Sequence[str], Type[PythonEnum]]] = None, ): if name in labelnames: raise ValueError(f'Overlapping labels for Enum metric: {name}') - if not states: + self._enum_type: Optional[Type[PythonEnum]] = None + if isinstance(states, type): + self._enum_type = states + state_names: Optional[Sequence[str]] = [member.name for member in states] + else: + state_names = states + if not state_names: raise ValueError(f'No states provided for Enum metric: {name}') super().__init__( name=name, @@ -790,15 +799,20 @@ def __init__(self, registry=registry, _labelvalues=_labelvalues, ) - self._kwargs['states'] = self._states = states + self._kwargs['states'] = states + self._states = state_names def _metric_init(self) -> None: self._value = 0 self._lock = Lock() - def state(self, state: str) -> None: + def state(self, state: Union[str, PythonEnum]) -> None: """Set enum metric state.""" self._raise_if_not_observable() + if self._enum_type is not None and isinstance(state, PythonEnum): + if not isinstance(state, self._enum_type): + raise ValueError(f'Invalid enum type for Enum metric: {self._name}') + state = state.name with self._lock: self._value = self._states.index(state) diff --git a/tests/test_core.py b/tests/test_core.py index 2c29f79a..b613c5eb 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1,4 +1,6 @@ from concurrent.futures import ThreadPoolExecutor +from enum import Enum as PythonEnum +from enum import IntEnum import os from threading import Barrier, Event import time @@ -14,6 +16,7 @@ StateSetMetricFamily, Summary, SummaryMetricFamily, UntypedMetricFamily, ) from prometheus_client.decorator import getargspec +from prometheus_client.exposition import generate_latest from prometheus_client.metrics import _get_use_created from prometheus_client.validation import ( disable_legacy_validation, enable_legacy_validation, @@ -592,6 +595,62 @@ def test_labels(self): self.assertEqual(1, self.registry.get_sample_value('el', {'l': 'a', 'el': 'c'})) self.assertRaises(ValueError, self.labels.state, 'a') + def test_python_enum(self): + class State(PythonEnum): + starting = 10 + running = 20 + stopped = 30 + finished = 30 + + metric = Enum('task_state', 'help', states=State, registry=self.registry) + self.assertEqual(1, self.registry.get_sample_value('task_state', {'task_state': 'starting'})) + metric.state(State.running) + self.assertEqual(0, self.registry.get_sample_value('task_state', {'task_state': 'starting'})) + self.assertEqual(1, self.registry.get_sample_value('task_state', {'task_state': 'running'})) + self.assertIn(b'task_state{task_state="running"} 1.0\n', generate_latest(self.registry)) + metric.state('stopped') + metric.state(State.finished) + self.assertEqual(1, self.registry.get_sample_value('task_state', {'task_state': 'stopped'})) + self.assertEqual(3, len(metric.collect()[0].samples)) + + def test_python_enum_labels(self): + class State(PythonEnum): + idle = 'Idle worker' + busy = 'Busy worker' + + metric = Enum('worker_state', 'help', ['worker'], states=State, registry=self.registry) + metric.labels('first').state(State.busy) + metric.labels('second') + self.assertEqual(1, self.registry.get_sample_value('worker_state', {'worker': 'first', 'worker_state': 'busy'})) + self.assertEqual(0, self.registry.get_sample_value('worker_state', {'worker': 'first', 'worker_state': 'idle'})) + self.assertEqual(1, self.registry.get_sample_value('worker_state', {'worker': 'second', 'worker_state': 'idle'})) + self.assertRaises(ValueError, metric.state, State.busy) + + def test_python_enum_rejects_unknown_state(self): + class State(IntEnum): + idle = 1 + busy = 2 + + class OtherState(IntEnum): + idle = 1 + busy = 2 + + metric = Enum('task_state', 'help', states=State, registry=self.registry) + metric.state(State.busy) + for state in (OtherState.idle, 'missing', 1): + with self.subTest(state=state): + self.assertRaises(ValueError, metric.state, state) + self.assertEqual(1, self.registry.get_sample_value('task_state', {'task_state': 'busy'})) + + def test_empty_python_enum_does_not_pollute_registry(self): + class EmptyState(PythonEnum): + pass + + with self.assertRaisesRegex(ValueError, 'No states provided'): + Enum('task_state', 'help', states=EmptyState, registry=self.registry) + Enum('task_state', 'help', states=['idle'], registry=self.registry) + self.assertEqual(1, self.registry.get_sample_value('task_state', {'task_state': 'idle'})) + def test_overlapping_labels(self): with pytest.raises(ValueError): Enum('e', 'help', registry=None, labelnames=['e'])