Use EntityDescription - sabnzbd (#55788)

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,
             }

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()