Use custom attributes description in Sonarr (#79845)

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
This commit is contained in:
Robert Hillis 2022-10-20 04:44:17 -04:00 committed by GitHub
parent 8a1cc05e0c
commit 41faa017cd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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: