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'])