diff --git a/homeassistant/components/sonarr/sensor.py b/homeassistant/components/sonarr/sensor.py index da0e4c5af8c..fe440b01be0 100644 --- a/homeassistant/components/sonarr/sensor.py +++ b/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: