Use enum for possible torrent statuses

This commit is contained in:
Joel Heaps 2023-09-28 22:24:49 -05:00
parent e8c0141365
commit 2f07f00592

View File

@ -9,6 +9,7 @@ from prometheus_client import start_http_server
from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily, REGISTRY from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily, REGISTRY
import logging import logging
from pythonjsonlogger import jsonlogger from pythonjsonlogger import jsonlogger
from enum import StrEnum, auto
# Enable dumps on stderr in case of segfault # Enable dumps on stderr in case of segfault
@ -16,16 +17,15 @@ faulthandler.enable()
logger = logging.getLogger() logger = logging.getLogger()
class QbittorrentMetricsCollector: class TorrentStatus(StrEnum):
TORRENT_STATUSES = [ CHECKING = auto()
"checking", COMPLETE = auto()
"complete", ERRORED = auto()
"downloading", PAUSED = auto()
"errored", UPLOADING = auto()
"paused",
"uploading",
]
class QbittorrentMetricsCollector:
def __init__(self, config): def __init__(self, config):
self.config = config self.config = config
self.client = Client( self.client = Client(
@ -125,8 +125,8 @@ class QbittorrentMetricsCollector:
or (category == "Uncategorized" and t["category"] == "") or (category == "Uncategorized" and t["category"] == "")
] ]
for status in self.TORRENT_STATUSES: for status in TorrentStatus:
status_prop = f"is_{status}" status_prop = f"is_{status.value}"
status_torrents = [ status_torrents = [
t t
for t in category_torrents for t in category_torrents
@ -139,12 +139,12 @@ class QbittorrentMetricsCollector:
"name": f"{self.config['metrics_prefix']}_torrents_count", "name": f"{self.config['metrics_prefix']}_torrents_count",
"value": len(status_torrents), "value": len(status_torrents),
"labels": { "labels": {
"status": status, "status": status.value,
"category": category, "category": category,
}, },
"help": ( "help": (
f"Number of torrents in status {status} under category" f"Number of torrents in status {status.value} under"
f" {category}" f" category {category}"
), ),
} }
) )