Use EntityDescription - radarr (#54997)

homeassistant/components/radarr/sensor.py

@@ -1,12 +1,19 @@
 """Support for Radarr."""
+from __future__ import annotations
+
 from datetime import datetime, timedelta
 import logging
 import time
+from typing import Any
 
 import requests
 import voluptuous as vol
 
-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_HOST,
@@ -42,14 +49,46 @@ DEFAULT_UNIT = DATA_GIGABYTES
 
 SCAN_INTERVAL = timedelta(minutes=10)
 
-SENSOR_TYPES = {
+SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
-    "diskspace": ["Disk Space", DATA_GIGABYTES, "mdi:harddisk"],
+    SensorEntityDescription(
-    "upcoming": ["Upcoming", "Movies", "mdi:television"],
+        key="diskspace",
-    "wanted": ["Wanted", "Movies", "mdi:television"],
+        name="Disk Space",
-    "movies": ["Movies", "Movies", "mdi:television"],
+        native_unit_of_measurement=DATA_GIGABYTES,
-    "commands": ["Commands", "Commands", "mdi:code-braces"],
+        icon="mdi:harddisk",
-    "status": ["Status", "Status", "mdi:information"],
+    ),
-}
+    SensorEntityDescription(
+        key="upcoming",
+        name="Upcoming",
+        native_unit_of_measurement="Movies",
+        icon="mdi:television",
+    ),
+    SensorEntityDescription(
+        key="wanted",
+        name="Wanted",
+        native_unit_of_measurement="Movies",
+        icon="mdi:television",
+    ),
+    SensorEntityDescription(
+        key="movies",
+        name="Movies",
+        native_unit_of_measurement="Movies",
+        icon="mdi:television",
+    ),
+    SensorEntityDescription(
+        key="commands",
+        name="Commands",
+        native_unit_of_measurement="Commands",
+        icon="mdi:code-braces",
+    ),
+    SensorEntityDescription(
+        key="status",
+        name="Status",
+        native_unit_of_measurement="Status",
+        icon="mdi:information",
+    ),
+)
+
+SENSOR_KEYS: list[str] = [desc.key for desc in SENSOR_TYPES]
 
 ENDPOINTS = {
     "diskspace": "{0}://{1}:{2}/{3}api/diskspace",
@@ -78,7 +117,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
         vol.Optional(CONF_HOST, default=DEFAULT_HOST): cv.string,
         vol.Optional(CONF_INCLUDED, default=[]): cv.ensure_list,
         vol.Optional(CONF_MONITORED_CONDITIONS, default=["movies"]): vol.All(
-            cv.ensure_list, [vol.In(list(SENSOR_TYPES))]
+            cv.ensure_list, [vol.In(SENSOR_KEYS)]
         ),
         vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
         vol.Optional(CONF_SSL, default=False): cv.boolean,
@@ -90,15 +129,21 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
 
 def setup_platform(hass, config, add_entities, discovery_info=None):
     """Set up the Radarr platform."""
-    conditions = config.get(CONF_MONITORED_CONDITIONS)
+    conditions = config[CONF_MONITORED_CONDITIONS]
-    add_entities([RadarrSensor(hass, config, sensor) for sensor in conditions], True)
+    entities = [
+        RadarrSensor(hass, config, description)
+        for description in SENSOR_TYPES
+        if description.key in conditions
+    ]
+    add_entities(entities, True)
 
 
 class RadarrSensor(SensorEntity):
     """Implementation of the Radarr sensor."""
 
-    def __init__(self, hass, conf, sensor_type):
+    def __init__(self, hass, conf, description: SensorEntityDescription):
         """Create Radarr entity."""
+        self.entity_description = description
 
         self.conf = conf
         self.host = conf.get(CONF_HOST)
@@ -110,78 +155,55 @@ class RadarrSensor(SensorEntity):
         self.included = conf.get(CONF_INCLUDED)
         self.days = int(conf.get(CONF_DAYS))
         self.ssl = "https" if conf.get(CONF_SSL) else "http"
-        self._state = None
+        self.data: list[Any] = []
-        self.data = []
+        self._attr_name = f"Radarr {description.name}"
-        self.type = sensor_type
+        if description.key == "diskspace":
-        self._name = SENSOR_TYPES[self.type][0]
+            self._attr_native_unit_of_measurement = conf.get(CONF_UNIT)
-        if self.type == "diskspace":
+        self._attr_available = False
-            self._unit = conf.get(CONF_UNIT)
-        else:
-            self._unit = SENSOR_TYPES[self.type][1]
-        self._icon = SENSOR_TYPES[self.type][2]
-        self._available = False
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return "{} {}".format("Radarr", self._name)
-
-    @property
-    def native_value(self):
-        """Return sensor state."""
-        return self._state
-
-    @property
-    def available(self):
-        """Return sensor availability."""
-        return self._available
-
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit of the sensor."""
-        return self._unit
 
     @property
     def extra_state_attributes(self):
         """Return the state attributes of the sensor."""
         attributes = {}
-        if self.type == "upcoming":
+        sensor_type = self.entity_description.key
+        if sensor_type == "upcoming":
             for movie in self.data:
                 attributes[to_key(movie)] = get_release_date(movie)
-        elif self.type == "commands":
+        elif sensor_type == "commands":
             for command in self.data:
                 attributes[command["name"]] = command["state"]
-        elif self.type == "diskspace":
+        elif sensor_type == "diskspace":
             for data in self.data:
-                free_space = to_unit(data["freeSpace"], self._unit)
+                free_space = to_unit(data["freeSpace"], self.native_unit_of_measurement)
-                total_space = to_unit(data["totalSpace"], self._unit)
+                total_space = to_unit(
+                    data["totalSpace"], self.native_unit_of_measurement
+                )
                 percentage_used = (
                     0 if total_space == 0 else free_space / total_space * 100
                 )
                 attributes[data["path"]] = "{:.2f}/{:.2f}{} ({:.2f}%)".format(
-                    free_space, total_space, self._unit, percentage_used
+                    free_space,
+                    total_space,
+                    self.native_unit_of_measurement,
+                    percentage_used,
                 )
-        elif self.type == "movies":
+        elif sensor_type == "movies":
             for movie in self.data:
                 attributes[to_key(movie)] = movie["downloaded"]
-        elif self.type == "status":
+        elif sensor_type == "status":
             attributes = self.data
 
         return attributes
 
-    @property
-    def icon(self):
-        """Return the icon of the sensor."""
-        return self._icon
-
     def update(self):
         """Update the data for the sensor."""
+        sensor_type = self.entity_description.key
         time_zone = dt_util.get_time_zone(self.hass.config.time_zone)
         start = get_date(time_zone)
         end = get_date(time_zone, self.days)
         try:
             res = requests.get(
-                ENDPOINTS[self.type].format(
+                ENDPOINTS[sensor_type].format(
                     self.ssl, self.host, self.port, self.urlbase, start, end
                 ),
                 headers={"X-Api-Key": self.apikey},
@@ -189,15 +211,15 @@ class RadarrSensor(SensorEntity):
             )
         except OSError:
             _LOGGER.warning("Host %s is not available", self.host)
-            self._available = False
+            self._attr_available = False
-            self._state = None
+            self._attr_native_value = None
             return
 
         if res.status_code == HTTP_OK:
-            if self.type in ("upcoming", "movies", "commands"):
+            if sensor_type in ("upcoming", "movies", "commands"):
                 self.data = res.json()
-                self._state = len(self.data)
+                self._attr_native_value = len(self.data)
-            elif self.type == "diskspace":
+            elif sensor_type == "diskspace":
                 # If included paths are not provided, use all data
                 if self.included == []:
                     self.data = res.json()
@@ -206,13 +228,16 @@ class RadarrSensor(SensorEntity):
                     self.data = list(
                         filter(lambda x: x["path"] in self.included, res.json())
                     )
-                self._state = "{:.2f}".format(
+                self._attr_native_value = "{:.2f}".format(
-                    to_unit(sum(data["freeSpace"] for data in self.data), self._unit)
+                    to_unit(
+                        sum(data["freeSpace"] for data in self.data),
+                        self.native_unit_of_measurement,
+                    )
                 )
-            elif self.type == "status":
+            elif sensor_type == "status":
                 self.data = res.json()
-                self._state = self.data["version"]
+                self._attr_native_value = self.data["version"]
-            self._available = True
+            self._attr_available = True
 
 
 def get_date(zone, offset=0):