Use custom attributes description in Sonarr (#79845)

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>

homeassistant/components/sonarr/sensor.py

@@ -6,6 +6,7 @@ from dataclasses import dataclass
 from typing import Any, Generic
 
 from aiopyarr import (
+    Command,
     Diskspace,
     SonarrCalendar,
     SonarrQueue,
@@ -30,6 +31,7 @@ from .entity import SonarrEntity
 class SonarrSensorEntityDescriptionMixIn(Generic[SonarrDataT]):
     """Mixin for Sonarr sensor."""
 
+    attributes_fn: Callable[[SonarrDataT], dict[str, str]]
     value_fn: Callable[[SonarrDataT], StateType]
 
 
@@ -40,14 +42,52 @@ class SonarrSensorEntityDescription(
     """Class to describe a Sonarr sensor."""
 
 
+def get_disk_space_attr(disks: list[Diskspace]) -> dict[str, str]:
+    """Create the attributes for disk space."""
+    attrs: dict[str, str] = {}
+    for disk in disks:
+        free = disk.freeSpace / 1024**3
+        total = disk.totalSpace / 1024**3
+        usage = free / total * 100
+        attrs[disk.path] = f"{free:.2f}/{total:.2f}{DATA_GIGABYTES} ({usage:.2f}%)"
+    return attrs
+
+
+def get_queue_attr(queue: SonarrQueue) -> dict[str, str]:
+    """Create the attributes for series queue."""
+    attrs: dict[str, str] = {}
+    for item in queue.records:
+        remaining = 1 if item.size == 0 else item.sizeleft / item.size
+        remaining_pct = 100 * (1 - remaining)
+        identifier = (
+            f"S{item.episode.seasonNumber:02d}E{item.episode. episodeNumber:02d}"
+        )
+        attrs[f"{item.series.title} {identifier}"] = f"{remaining_pct:.2f}%"
+    return attrs
+
+
+def get_wanted_attr(wanted: SonarrWantedMissing) -> dict[str, str]:
+    """Create the attributes for missing series."""
+    attrs: dict[str, str] = {}
+    for item in wanted.records:
+        identifier = f"S{item.seasonNumber:02d}E{item.episodeNumber:02d}"
+
+        name = f"{item.series.title} {identifier}"
+        attrs[name] = dt_util.as_local(
+            item.airDateUtc.replace(tzinfo=dt_util.UTC)
+        ).isoformat()
+    return attrs
+
+
 SENSOR_TYPES: dict[str, SonarrSensorEntityDescription[Any]] = {
-    "commands": SonarrSensorEntityDescription(
+    "commands": SonarrSensorEntityDescription[list[Command]](
         key="commands",
         name="Commands",
         icon="mdi:code-braces",
         native_unit_of_measurement="Commands",
         entity_registry_enabled_default=False,
         value_fn=len,
+        attributes_fn=lambda data: {c.name: c.status for c in data},
     ),
     "diskspace": SonarrSensorEntityDescription[list[Diskspace]](
         key="diskspace",
@@ -56,6 +96,7 @@ SENSOR_TYPES: dict[str, SonarrSensorEntityDescription[Any]] = {
         native_unit_of_measurement=DATA_GIGABYTES,
         entity_registry_enabled_default=False,
         value_fn=lambda data: f"{sum(disk.freeSpace for disk in data) / 1024**3:.2f}",
+        attributes_fn=get_disk_space_attr,
     ),
     "queue": SonarrSensorEntityDescription[SonarrQueue](
         key="queue",
@@ -64,6 +105,7 @@ SENSOR_TYPES: dict[str, SonarrSensorEntityDescription[Any]] = {
         native_unit_of_measurement="Episodes",
         entity_registry_enabled_default=False,
         value_fn=lambda data: data.totalRecords,
+        attributes_fn=get_queue_attr,
     ),
     "series": SonarrSensorEntityDescription[list[SonarrSeries]](
         key="series",
@@ -72,6 +114,10 @@ SENSOR_TYPES: dict[str, SonarrSensorEntityDescription[Any]] = {
         native_unit_of_measurement="Series",
         entity_registry_enabled_default=False,
         value_fn=len,
+        attributes_fn=lambda data: {
+            i.title: f"{getattr(i.statistics,'episodeFileCount', 0)}/{getattr(i.statistics, 'episodeCount', 0)} Episodes"
+            for i in data
+        },
     ),
     "upcoming": SonarrSensorEntityDescription[list[SonarrCalendar]](
         key="upcoming",
@@ -79,6 +125,9 @@ SENSOR_TYPES: dict[str, SonarrSensorEntityDescription[Any]] = {
         icon="mdi:television",
         native_unit_of_measurement="Episodes",
         value_fn=len,
+        attributes_fn=lambda data: {
+            e.series.title: f"S{e.seasonNumber:02d}E{e.episodeNumber:02d}" for e in data
+        },
     ),
     "wanted": SonarrSensorEntityDescription[SonarrWantedMissing](
         key="wanted",
@@ -87,6 +136,7 @@ SENSOR_TYPES: dict[str, SonarrSensorEntityDescription[Any]] = {
         native_unit_of_measurement="Episodes",
         entity_registry_enabled_default=False,
         value_fn=lambda data: data.totalRecords,
+        attributes_fn=get_wanted_attr,
     ),
 }
 
@@ -109,56 +159,13 @@ async def async_setup_entry(
 class SonarrSensor(SonarrEntity[SonarrDataT], SensorEntity):
     """Implementation of the Sonarr sensor."""
 
-    coordinator: SonarrDataUpdateCoordinator
+    coordinator: SonarrDataUpdateCoordinator[SonarrDataT]
     entity_description: SonarrSensorEntityDescription[SonarrDataT]
 
     @property
-    def extra_state_attributes(self) -> dict[str, str] | None:
+    def extra_state_attributes(self) -> dict[str, str]:
         """Return the state attributes of the entity."""
-        attrs = {}
-        key = self.entity_description.key
-        data = self.coordinator.data
-
-        if key == "diskspace":
-            for disk in data:
-                free = disk.freeSpace / 1024**3
-                total = disk.totalSpace / 1024**3
-                usage = free / total * 100
-
-                attrs[
-                    disk.path
-                ] = f"{free:.2f}/{total:.2f}{self.unit_of_measurement} ({usage:.2f}%)"
-        elif key == "commands":
-            for command in data:
-                attrs[command.name] = command.status
-        elif key == "queue":
-            for item in data.records:
-                remaining = 1 if item.size == 0 else item.sizeleft / item.size
-                remaining_pct = 100 * (1 - remaining)
-                identifier = f"S{item.episode.seasonNumber:02d}E{item.episode. episodeNumber:02d}"
-
-                name = f"{item.series.title} {identifier}"
-                attrs[name] = f"{remaining_pct:.2f}%"
-        elif key == "series":
-            for item in data:
-                stats = item.statistics
-                attrs[
-                    item.title
-                ] = f"{getattr(stats,'episodeFileCount', 0)}/{getattr(stats, 'episodeCount', 0)} Episodes"
-        elif key == "upcoming":
-            for episode in data:
-                identifier = f"S{episode.seasonNumber:02d}E{episode.episodeNumber:02d}"
-                attrs[episode.series.title] = identifier
-        elif key == "wanted":
-            for item in data.records:
-                identifier = f"S{item.seasonNumber:02d}E{item.episodeNumber:02d}"
-
-                name = f"{item.series.title} {identifier}"
-                attrs[name] = dt_util.as_local(
-                    item.airDateUtc.replace(tzinfo=dt_util.UTC)
-                ).isoformat()
-
-        return attrs
+        return self.entity_description.attributes_fn(self.coordinator.data)
 
     @property
     def native_value(self) -> StateType: