Use EntityDescription - startca (#55036)

homeassistant/components/startca/sensor.py

@@ -1,4 +1,6 @@
 """Support for Start.ca Bandwidth Monitor."""
+from __future__ import annotations
+
 from datetime import timedelta
 import logging
 from xml.parsers.expat import ExpatError
@@ -7,7 +9,11 @@ import async_timeout
 import voluptuous as vol
 import xmltodict
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    SensorEntity,
+    SensorEntityDescription,
+)
 from homeassistant.const import (
     CONF_API_KEY,
     CONF_MONITORED_VARIABLES,
@@ -28,25 +34,87 @@ CONF_TOTAL_BANDWIDTH = "total_bandwidth"
 MIN_TIME_BETWEEN_UPDATES = timedelta(hours=1)
 REQUEST_TIMEOUT = 5  # seconds
 
-SENSOR_TYPES = {
-    "usage": ["Usage Ratio", PERCENTAGE, "mdi:percent"],
-    "usage_gb": ["Usage", DATA_GIGABYTES, "mdi:download"],
-    "limit": ["Data limit", DATA_GIGABYTES, "mdi:download"],
-    "used_download": ["Used Download", DATA_GIGABYTES, "mdi:download"],
-    "used_upload": ["Used Upload", DATA_GIGABYTES, "mdi:upload"],
-    "used_total": ["Used Total", DATA_GIGABYTES, "mdi:download"],
-    "grace_download": ["Grace Download", DATA_GIGABYTES, "mdi:download"],
-    "grace_upload": ["Grace Upload", DATA_GIGABYTES, "mdi:upload"],
-    "grace_total": ["Grace Total", DATA_GIGABYTES, "mdi:download"],
-    "total_download": ["Total Download", DATA_GIGABYTES, "mdi:download"],
-    "total_upload": ["Total Upload", DATA_GIGABYTES, "mdi:download"],
-    "used_remaining": ["Remaining", DATA_GIGABYTES, "mdi:download"],
-}
+SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key="usage",
+        name="Usage Ratio",
+        native_unit_of_measurement=PERCENTAGE,
+        icon="mdi:percent",
+    ),
+    SensorEntityDescription(
+        key="usage_gb",
+        name="Usage",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+    SensorEntityDescription(
+        key="limit",
+        name="Data limit",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+    SensorEntityDescription(
+        key="used_download",
+        name="Used Download",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+    SensorEntityDescription(
+        key="used_upload",
+        name="Used Upload",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:upload",
+    ),
+    SensorEntityDescription(
+        key="used_total",
+        name="Used Total",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+    SensorEntityDescription(
+        key="grace_download",
+        name="Grace Download",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+    SensorEntityDescription(
+        key="grace_upload",
+        name="Grace Upload",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:upload",
+    ),
+    SensorEntityDescription(
+        key="grace_total",
+        name="Grace Total",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+    SensorEntityDescription(
+        key="total_download",
+        name="Total Download",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+    SensorEntityDescription(
+        key="total_upload",
+        name="Total Upload",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+    SensorEntityDescription(
+        key="used_remaining",
+        name="Remaining",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:download",
+    ),
+)
+
+SENSOR_KEYS: list[str] = [desc.key for desc in SENSOR_TYPES]
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Required(CONF_MONITORED_VARIABLES): vol.All(
-            cv.ensure_list, [vol.In(SENSOR_TYPES)]
+            cv.ensure_list, [vol.In(SENSOR_KEYS)]
         ),
         vol.Required(CONF_API_KEY): cv.string,
         vol.Required(CONF_TOTAL_BANDWIDTH): cv.positive_int,
@@ -58,8 +126,8 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
 async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
     """Set up the sensor platform."""
     websession = async_get_clientsession(hass)
-    apikey = config.get(CONF_API_KEY)
-    bandwidthcap = config.get(CONF_TOTAL_BANDWIDTH)
+    apikey = config[CONF_API_KEY]
+    bandwidthcap = config[CONF_TOTAL_BANDWIDTH]
 
     ts_data = StartcaData(hass.loop, websession, apikey, bandwidthcap)
     ret = await ts_data.async_update()
@@ -67,51 +135,32 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
         _LOGGER.error("Invalid Start.ca API key: %s", apikey)
         return
 
-    name = config.get(CONF_NAME)
-    sensors = []
-    for variable in config[CONF_MONITORED_VARIABLES]:
-        sensors.append(StartcaSensor(ts_data, variable, name))
-    async_add_entities(sensors, True)
+    name = config[CONF_NAME]
+    monitored_variables = config[CONF_MONITORED_VARIABLES]
+    entities = [
+        StartcaSensor(ts_data, name, description)
+        for description in SENSOR_TYPES
+        if description.key in monitored_variables
+    ]
+    async_add_entities(entities, True)
 
 
 class StartcaSensor(SensorEntity):
     """Representation of Start.ca Bandwidth sensor."""
 
-    def __init__(self, startcadata, sensor_type, name):
+    def __init__(self, startcadata, name, description: SensorEntityDescription):
         """Initialize the sensor."""
-        self.client_name = name
-        self.type = sensor_type
-        self._name = SENSOR_TYPES[sensor_type][0]
-        self._unit_of_measurement = SENSOR_TYPES[sensor_type][1]
-        self._icon = SENSOR_TYPES[sensor_type][2]
+        self.entity_description = description
         self.startcadata = startcadata
-        self._state = None
 
-    @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 native_unit_of_measurement(self):
-        """Return the unit of measurement of this entity, if any."""
-        return self._unit_of_measurement
-
-    @property
-    def icon(self):
-        """Icon to use in the frontend, if any."""
-        return self._icon
+        self._attr_name = f"{name} {description.name}"
 
     async def async_update(self):
         """Get the latest data from Start.ca and update the state."""
         await self.startcadata.async_update()
-        if self.type in self.startcadata.data:
-            self._state = round(self.startcadata.data[self.type], 2)
+        sensor_type = self.entity_description.key
+        if sensor_type in self.startcadata.data:
+            self._attr_native_value = round(self.startcadata.data[sensor_type], 2)
 
 
 class StartcaData: