diff --git a/homeassistant/components/sabnzbd/__init__.py b/homeassistant/components/sabnzbd/__init__.py index a420ca53814..f6e15bd074c 100644 --- a/homeassistant/components/sabnzbd/__init__.py +++ b/homeassistant/components/sabnzbd/__init__.py @@ -1,6 +1,7 @@ """Support for monitoring an SABnzbd NZB client.""" from __future__ import annotations +from dataclasses import dataclass from datetime import timedelta import logging @@ -8,6 +9,7 @@ from pysabnzbd import SabnzbdApi, SabnzbdApiException import voluptuous as vol from homeassistant.components.discovery import SERVICE_SABNZBD +from homeassistant.components.sensor import SensorEntityDescription from homeassistant.const import ( CONF_API_KEY, CONF_HOST, @@ -52,19 +54,87 @@ SERVICE_SET_SPEED = "set_speed" SIGNAL_SABNZBD_UPDATED = "sabnzbd_updated" -SENSOR_TYPES = { - "current_status": ["Status", None, "status"], - "speed": ["Speed", DATA_RATE_MEGABYTES_PER_SECOND, "kbpersec"], - "queue_size": ["Queue", DATA_MEGABYTES, "mb"], - "queue_remaining": ["Left", DATA_MEGABYTES, "mbleft"], - "disk_size": ["Disk", DATA_GIGABYTES, "diskspacetotal1"], - "disk_free": ["Disk Free", DATA_GIGABYTES, "diskspace1"], - "queue_count": ["Queue Count", None, "noofslots_total"], - "day_size": ["Daily Total", DATA_GIGABYTES, "day_size"], - "week_size": ["Weekly Total", DATA_GIGABYTES, "week_size"], - "month_size": ["Monthly Total", DATA_GIGABYTES, "month_size"], - "total_size": ["Total", DATA_GIGABYTES, "total_size"], -} + +@dataclass +class SabnzbdRequiredKeysMixin: + """Mixin for required keys.""" + + field_name: str + + +@dataclass +class SabnzbdSensorEntityDescription(SensorEntityDescription, SabnzbdRequiredKeysMixin): + """Describes Sabnzbd sensor entity.""" + + +SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = ( + SabnzbdSensorEntityDescription( + key="current_status", + name="Status", + field_name="status", + ), + SabnzbdSensorEntityDescription( + key="speed", + name="Speed", + native_unit_of_measurement=DATA_RATE_MEGABYTES_PER_SECOND, + field_name="kbpersec", + ), + SabnzbdSensorEntityDescription( + key="queue_size", + name="Queue", + native_unit_of_measurement=DATA_MEGABYTES, + field_name="mb", + ), + SabnzbdSensorEntityDescription( + key="queue_remaining", + name="Left", + native_unit_of_measurement=DATA_MEGABYTES, + field_name="mbleft", + ), + SabnzbdSensorEntityDescription( + key="disk_size", + name="Disk", + native_unit_of_measurement=DATA_GIGABYTES, + field_name="diskspacetotal1", + ), + SabnzbdSensorEntityDescription( + key="disk_free", + name="Disk Free", + native_unit_of_measurement=DATA_GIGABYTES, + field_name="diskspace1", + ), + SabnzbdSensorEntityDescription( + key="queue_count", + name="Queue Count", + field_name="noofslots_total", + ), + SabnzbdSensorEntityDescription( + key="day_size", + name="Daily Total", + native_unit_of_measurement=DATA_GIGABYTES, + field_name="day_size", + ), + SabnzbdSensorEntityDescription( + key="week_size", + name="Weekly Total", + native_unit_of_measurement=DATA_GIGABYTES, + field_name="week_size", + ), + SabnzbdSensorEntityDescription( + key="month_size", + name="Monthly Total", + native_unit_of_measurement=DATA_GIGABYTES, + field_name="month_size", + ), + SabnzbdSensorEntityDescription( + key="total_size", + name="Total", + native_unit_of_measurement=DATA_GIGABYTES, + field_name="total_size", + ), +) + +SENSOR_KEYS: list[str] = [desc.key for desc in SENSOR_TYPES] SPEED_LIMIT_SCHEMA = vol.Schema( {vol.Optional(ATTR_SPEED, default=DEFAULT_SPEED_LIMIT): cv.string} @@ -80,7 +150,7 @@ CONFIG_SCHEMA = vol.Schema( vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string, vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port, vol.Optional(CONF_SENSORS): vol.All( - cv.ensure_list, [vol.In(SENSOR_TYPES)] + cv.ensure_list, [vol.In(SENSOR_KEYS)] ), vol.Optional(CONF_SSL, default=DEFAULT_SSL): cv.boolean, } diff --git a/homeassistant/components/sabnzbd/sensor.py b/homeassistant/components/sabnzbd/sensor.py index ffe57e608bf..3079bd75601 100644 --- a/homeassistant/components/sabnzbd/sensor.py +++ b/homeassistant/components/sabnzbd/sensor.py @@ -2,7 +2,12 @@ from homeassistant.components.sensor import SensorEntity from homeassistant.helpers.dispatcher import async_dispatcher_connect -from . import DATA_SABNZBD, SENSOR_TYPES, SIGNAL_SABNZBD_UPDATED +from . import ( + DATA_SABNZBD, + SENSOR_TYPES, + SIGNAL_SABNZBD_UPDATED, + SabnzbdSensorEntityDescription, +) async def async_setup_platform(hass, config, async_add_entities, discovery_info=None): @@ -14,22 +19,27 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info= sensors = sab_api_data.sensors client_name = sab_api_data.name async_add_entities( - [SabnzbdSensor(sensor, sab_api_data, client_name) for sensor in sensors] + [ + SabnzbdSensor(sab_api_data, client_name, description) + for description in SENSOR_TYPES + if description.key in sensors + ] ) class SabnzbdSensor(SensorEntity): """Representation of an SABnzbd sensor.""" - def __init__(self, sensor_type, sabnzbd_api_data, client_name): + entity_description: SabnzbdSensorEntityDescription + _attr_should_poll = False + + def __init__( + self, sabnzbd_api_data, client_name, description: SabnzbdSensorEntityDescription + ): """Initialize the sensor.""" - self._client_name = client_name - self._field_name = SENSOR_TYPES[sensor_type][2] - self._name = SENSOR_TYPES[sensor_type][0] + self.entity_description = description self._sabnzbd_api = sabnzbd_api_data - self._state = None - self._type = sensor_type - self._unit_of_measurement = SENSOR_TYPES[sensor_type][1] + self._attr_name = f"{client_name} {description.name}" async def async_added_to_hass(self): """Call when entity about to be added to hass.""" @@ -39,33 +49,15 @@ class SabnzbdSensor(SensorEntity): ) ) - @property - def name(self): - """Return the name of the sensor.""" - return f"{self._client_name} {self._name}" - - @property - def native_value(self): - """Return the state of the sensor.""" - return self._state - - @property - def should_poll(self): - """Don't poll. Will be updated by dispatcher signal.""" - return False - - @property - def native_unit_of_measurement(self): - """Return the unit of measurement of this entity, if any.""" - return self._unit_of_measurement - def update_state(self, args): """Get the latest data and updates the states.""" - self._state = self._sabnzbd_api.get_queue_field(self._field_name) + self._attr_native_value = self._sabnzbd_api.get_queue_field( + self.entity_description.field_name + ) - if self._type == "speed": - self._state = round(float(self._state) / 1024, 1) - elif "size" in self._type: - self._state = round(float(self._state), 2) + if self.entity_description.key == "speed": + self._attr_native_value = round(float(self._attr_native_value) / 1024, 1) + elif "size" in self.entity_description.key: + self._attr_native_value = round(float(self._attr_native_value), 2) self.schedule_update_ha_state()