Use entity descriptions for sonarr (#55818)

homeassistant/components/sonarr/sensor.py

@@ -3,22 +3,16 @@ from __future__ import annotations
 
 from datetime import timedelta
 import logging
+from typing import Any
 
 from sonarr import Sonarr, SonarrConnectionError, SonarrError
-from sonarr.models import (
-    CommandItem,
-    Disk,
-    Episode,
-    QueueItem,
-    SeriesItem,
-    WantedResults,
-)
 
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import DATA_GIGABYTES
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import StateType
 import homeassistant.util.dt as dt_util
 
 from .const import CONF_UPCOMING_DAYS, CONF_WANTED_MAX_ITEMS, DATA_SONARR, DOMAIN
@@ -26,6 +20,50 @@ from .entity import SonarrEntity
 
 _LOGGER = logging.getLogger(__name__)
 
+SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key="commands",
+        name="Sonarr Commands",
+        icon="mdi:code-braces",
+        native_unit_of_measurement="Commands",
+        entity_registry_enabled_default=False,
+    ),
+    SensorEntityDescription(
+        key="diskspace",
+        name="Sonarr Disk Space",
+        icon="mdi:harddisk",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        entity_registry_enabled_default=False,
+    ),
+    SensorEntityDescription(
+        key="queue",
+        name="Sonarr Queue",
+        icon="mdi:download",
+        native_unit_of_measurement="Episodes",
+        entity_registry_enabled_default=False,
+    ),
+    SensorEntityDescription(
+        key="series",
+        name="Sonarr Shows",
+        icon="mdi:television",
+        native_unit_of_measurement="Series",
+        entity_registry_enabled_default=False,
+    ),
+    SensorEntityDescription(
+        key="upcoming",
+        name="Sonarr Upcoming",
+        icon="mdi:television",
+        native_unit_of_measurement="Episodes",
+    ),
+    SensorEntityDescription(
+        key="wanted",
+        name="Sonarr Wanted",
+        icon="mdi:television",
+        native_unit_of_measurement="Episodes",
+        entity_registry_enabled_default=False,
+    ),
+)
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -33,18 +71,12 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Sonarr sensors based on a config entry."""
-    options = entry.options
+    sonarr: Sonarr = hass.data[DOMAIN][entry.entry_id][DATA_SONARR]
-    sonarr = hass.data[DOMAIN][entry.entry_id][DATA_SONARR]
+    options: dict[str, Any] = dict(entry.options)
 
     entities = [
-        SonarrCommandsSensor(sonarr, entry.entry_id),
+        SonarrSensor(sonarr, entry.entry_id, description, options)
-        SonarrDiskspaceSensor(sonarr, entry.entry_id),
+        for description in SENSOR_TYPES
-        SonarrQueueSensor(sonarr, entry.entry_id),
-        SonarrSeriesSensor(sonarr, entry.entry_id),
-        SonarrUpcomingSensor(sonarr, entry.entry_id, days=options[CONF_UPCOMING_DAYS]),
-        SonarrWantedSensor(
-            sonarr, entry.entry_id, max_items=options[CONF_WANTED_MAX_ITEMS]
-        ),
     ]
 
     async_add_entities(entities, True)
@@ -78,23 +110,19 @@ class SonarrSensor(SonarrEntity, SensorEntity):
 
     def __init__(
         self,
-        *,
         sonarr: Sonarr,
         entry_id: str,
-        enabled_default: bool = True,
+        description: SensorEntityDescription,
-        icon: str,
+        options: dict[str, Any],
-        key: str,
-        name: str,
-        unit_of_measurement: str | None = None,
     ) -> None:
         """Initialize Sonarr sensor."""
-        self._key = key
+        self.entity_description = description
-        self._attr_name = name
+        self._attr_unique_id = f"{entry_id}_{description.key}"
-        self._attr_icon = icon
+
-        self._attr_unique_id = f"{entry_id}_{key}"
+        self.data: dict[str, Any] = {}
-        self._attr_native_unit_of_measurement = unit_of_measurement
+        self.last_update_success: bool = False
-        self._attr_entity_registry_enabled_default = enabled_default
+        self.upcoming_days: int = options[CONF_UPCOMING_DAYS]
-        self.last_update_success = False
+        self.wanted_max_items: int = options[CONF_WANTED_MAX_ITEMS]
 
         super().__init__(
             sonarr=sonarr,
@@ -107,76 +135,38 @@ class SonarrSensor(SonarrEntity, SensorEntity):
         """Return sensor availability."""
         return self.last_update_success
 
-
-class SonarrCommandsSensor(SonarrSensor):
-    """Defines a Sonarr Commands sensor."""
-
-    def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
-        """Initialize Sonarr Commands sensor."""
-        self._commands: list[CommandItem] = []
-
-        super().__init__(
-            sonarr=sonarr,
-            entry_id=entry_id,
-            icon="mdi:code-braces",
-            key="commands",
-            name=f"{sonarr.app.info.app_name} Commands",
-            unit_of_measurement="Commands",
-            enabled_default=False,
-        )
-
     @sonarr_exception_handler
     async def async_update(self) -> None:
         """Update entity."""
-        self._commands = await self.sonarr.commands()
+        key = self.entity_description.key
+
+        if key == "diskspace":
+            await self.sonarr.update()
+        elif key == "commands":
+            self.data[key] = await self.sonarr.commands()
+        elif key == "queue":
+            self.data[key] = await self.sonarr.queue()
+        elif key == "series":
+            self.data[key] = await self.sonarr.series()
+        elif key == "upcoming":
+            local = dt_util.start_of_local_day().replace(microsecond=0)
+            start = dt_util.as_utc(local)
+            end = start + timedelta(days=self.upcoming_days)
+
+            self.data[key] = await self.sonarr.calendar(
+                start=start.isoformat(), end=end.isoformat()
+            )
+        elif key == "wanted":
+            self.data[key] = await self.sonarr.wanted(page_size=self.wanted_max_items)
 
     @property
     def extra_state_attributes(self) -> dict[str, str] | None:
         """Return the state attributes of the entity."""
         attrs = {}
+        key = self.entity_description.key
 
-        for command in self._commands:
+        if key == "diskspace":
-            attrs[command.name] = command.state
+            for disk in self.sonarr.app.disks:
-
-        return attrs
-
-    @property
-    def native_value(self) -> int:
-        """Return the state of the sensor."""
-        return len(self._commands)
-
-
-class SonarrDiskspaceSensor(SonarrSensor):
-    """Defines a Sonarr Disk Space sensor."""
-
-    def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
-        """Initialize Sonarr Disk Space sensor."""
-        self._disks: list[Disk] = []
-        self._total_free = 0
-
-        super().__init__(
-            sonarr=sonarr,
-            entry_id=entry_id,
-            icon="mdi:harddisk",
-            key="diskspace",
-            name=f"{sonarr.app.info.app_name} Disk Space",
-            unit_of_measurement=DATA_GIGABYTES,
-            enabled_default=False,
-        )
-
-    @sonarr_exception_handler
-    async def async_update(self) -> None:
-        """Update entity."""
-        app = await self.sonarr.update()
-        self._disks = app.disks
-        self._total_free = sum(disk.free for disk in self._disks)
-
-    @property
-    def extra_state_attributes(self) -> dict[str, str] | None:
-        """Return the state attributes of the entity."""
-        attrs = {}
-
-        for disk in self._disks:
                 free = disk.free / 1024 ** 3
                 total = disk.total / 1024 ** 3
                 usage = free / total * 100
@@ -184,176 +174,51 @@ class SonarrDiskspaceSensor(SonarrSensor):
                 attrs[
                     disk.path
                 ] = f"{free:.2f}/{total:.2f}{self.unit_of_measurement} ({usage:.2f}%)"
-
+        elif key == "commands" and self.data.get(key) is not None:
-        return attrs
+            for command in self.data[key]:
-
+                attrs[command.name] = command.state
-    @property
+        elif key == "queue" and self.data.get(key) is not None:
-    def native_value(self) -> str:
+            for item in self.data[key]:
-        """Return the state of the sensor."""
-        free = self._total_free / 1024 ** 3
-        return f"{free:.2f}"
-
-
-class SonarrQueueSensor(SonarrSensor):
-    """Defines a Sonarr Queue sensor."""
-
-    def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
-        """Initialize Sonarr Queue sensor."""
-        self._queue: list[QueueItem] = []
-
-        super().__init__(
-            sonarr=sonarr,
-            entry_id=entry_id,
-            icon="mdi:download",
-            key="queue",
-            name=f"{sonarr.app.info.app_name} Queue",
-            unit_of_measurement="Episodes",
-            enabled_default=False,
-        )
-
-    @sonarr_exception_handler
-    async def async_update(self) -> None:
-        """Update entity."""
-        self._queue = await self.sonarr.queue()
-
-    @property
-    def extra_state_attributes(self) -> dict[str, str] | None:
-        """Return the state attributes of the entity."""
-        attrs = {}
-
-        for item in self._queue:
                 remaining = 1 if item.size == 0 else item.size_remaining / item.size
                 remaining_pct = 100 * (1 - remaining)
                 name = f"{item.episode.series.title} {item.episode.identifier}"
                 attrs[name] = f"{remaining_pct:.2f}%"
-
+        elif key == "series" and self.data.get(key) is not None:
-        return attrs
+            for item in self.data[key]:
-
-    @property
-    def native_value(self) -> int:
-        """Return the state of the sensor."""
-        return len(self._queue)
-
-
-class SonarrSeriesSensor(SonarrSensor):
-    """Defines a Sonarr Series sensor."""
-
-    def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
-        """Initialize Sonarr Series sensor."""
-        self._items: list[SeriesItem] = []
-
-        super().__init__(
-            sonarr=sonarr,
-            entry_id=entry_id,
-            icon="mdi:television",
-            key="series",
-            name=f"{sonarr.app.info.app_name} Shows",
-            unit_of_measurement="Series",
-            enabled_default=False,
-        )
-
-    @sonarr_exception_handler
-    async def async_update(self) -> None:
-        """Update entity."""
-        self._items = await self.sonarr.series()
-
-    @property
-    def extra_state_attributes(self) -> dict[str, str] | None:
-        """Return the state attributes of the entity."""
-        attrs = {}
-
-        for item in self._items:
                 attrs[item.series.title] = f"{item.downloaded}/{item.episodes} Episodes"
-
+        elif key == "upcoming" and self.data.get(key) is not None:
-        return attrs
+            for episode in self.data[key]:
-
-    @property
-    def native_value(self) -> int:
-        """Return the state of the sensor."""
-        return len(self._items)
-
-
-class SonarrUpcomingSensor(SonarrSensor):
-    """Defines a Sonarr Upcoming sensor."""
-
-    def __init__(self, sonarr: Sonarr, entry_id: str, days: int = 1) -> None:
-        """Initialize Sonarr Upcoming sensor."""
-        self._days = days
-        self._upcoming: list[Episode] = []
-
-        super().__init__(
-            sonarr=sonarr,
-            entry_id=entry_id,
-            icon="mdi:television",
-            key="upcoming",
-            name=f"{sonarr.app.info.app_name} Upcoming",
-            unit_of_measurement="Episodes",
-        )
-
-    @sonarr_exception_handler
-    async def async_update(self) -> None:
-        """Update entity."""
-        local = dt_util.start_of_local_day().replace(microsecond=0)
-        start = dt_util.as_utc(local)
-        end = start + timedelta(days=self._days)
-        self._upcoming = await self.sonarr.calendar(
-            start=start.isoformat(), end=end.isoformat()
-        )
-
-    @property
-    def extra_state_attributes(self) -> dict[str, str] | None:
-        """Return the state attributes of the entity."""
-        attrs = {}
-
-        for episode in self._upcoming:
                 attrs[episode.series.title] = episode.identifier
-
+        elif key == "wanted" and self.data.get(key) is not None:
-        return attrs
+            for episode in self.data[key].episodes:
-
-    @property
-    def native_value(self) -> int:
-        """Return the state of the sensor."""
-        return len(self._upcoming)
-
-
-class SonarrWantedSensor(SonarrSensor):
-    """Defines a Sonarr Wanted sensor."""
-
-    def __init__(self, sonarr: Sonarr, entry_id: str, max_items: int = 10) -> None:
-        """Initialize Sonarr Wanted sensor."""
-        self._max_items = max_items
-        self._results: WantedResults | None = None
-        self._total: int | None = None
-
-        super().__init__(
-            sonarr=sonarr,
-            entry_id=entry_id,
-            icon="mdi:television",
-            key="wanted",
-            name=f"{sonarr.app.info.app_name} Wanted",
-            unit_of_measurement="Episodes",
-            enabled_default=False,
-        )
-
-    @sonarr_exception_handler
-    async def async_update(self) -> None:
-        """Update entity."""
-        self._results = await self.sonarr.wanted(page_size=self._max_items)
-        self._total = self._results.total
-
-    @property
-    def extra_state_attributes(self) -> dict[str, str] | None:
-        """Return the state attributes of the entity."""
-        attrs: dict[str, str] = {}
-
-        if self._results is not None:
-            for episode in self._results.episodes:
                 name = f"{episode.series.title} {episode.identifier}"
                 attrs[name] = episode.airdate
 
         return attrs
 
     @property
-    def native_value(self) -> int | None:
+    def native_value(self) -> StateType:
         """Return the state of the sensor."""
-        return self._total
+        key = self.entity_description.key
+
+        if key == "diskspace":
+            total_free = sum(disk.free for disk in self.sonarr.app.disks)
+            free = total_free / 1024 ** 3
+            return f"{free:.2f}"
+
+        if key == "commands" and self.data.get(key) is not None:
+            return len(self.data[key])
+
+        if key == "queue" and self.data.get(key) is not None:
+            return len(self.data[key])
+
+        if key == "series" and self.data.get(key) is not None:
+            return len(self.data[key])
+
+        if key == "upcoming" and self.data.get(key) is not None:
+            return len(self.data[key])
+
+        if key == "wanted" and self.data.get(key) is not None:
+            return self.data[key].total
+
+        return None