core/homeassistant/components/sonarr/sensor.py

"""Support for Sonarr sensors."""
from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Generic

from aiopyarr import (
    Diskspace,
    SonarrCalendar,
    SonarrQueue,
    SonarrSeries,
    SonarrWantedMissing,
)

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 DOMAIN
from .coordinator import SonarrDataT, SonarrDataUpdateCoordinator
from .entity import SonarrEntity


@dataclass
class SonarrSensorEntityDescriptionMixIn(Generic[SonarrDataT]):
    """Mixin for Sonarr sensor."""

    value_fn: Callable[[SonarrDataT], StateType]


@dataclass
class SonarrSensorEntityDescription(
    SensorEntityDescription, SonarrSensorEntityDescriptionMixIn[SonarrDataT]
):
    """Class to describe a Sonarr sensor."""


SENSOR_TYPES: dict[str, SonarrSensorEntityDescription[Any]] = {
    "commands": SonarrSensorEntityDescription(
        key="commands",
        name="Commands",
        icon="mdi:code-braces",
        native_unit_of_measurement="Commands",
        entity_registry_enabled_default=False,
        value_fn=len,
    ),
    "diskspace": SonarrSensorEntityDescription[list[Diskspace]](
        key="diskspace",
        name="Disk space",
        icon="mdi:harddisk",
        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}",
    ),
    "queue": SonarrSensorEntityDescription[SonarrQueue](
        key="queue",
        name="Queue",
        icon="mdi:download",
        native_unit_of_measurement="Episodes",
        entity_registry_enabled_default=False,
        value_fn=lambda data: data.totalRecords,
    ),
    "series": SonarrSensorEntityDescription[list[SonarrSeries]](
        key="series",
        name="Shows",
        icon="mdi:television",
        native_unit_of_measurement="Series",
        entity_registry_enabled_default=False,
        value_fn=len,
    ),
    "upcoming": SonarrSensorEntityDescription[list[SonarrCalendar]](
        key="upcoming",
        name="Upcoming",
        icon="mdi:television",
        native_unit_of_measurement="Episodes",
        value_fn=len,
    ),
    "wanted": SonarrSensorEntityDescription[SonarrWantedMissing](
        key="wanted",
        name="Wanted",
        icon="mdi:television",
        native_unit_of_measurement="Episodes",
        entity_registry_enabled_default=False,
        value_fn=lambda data: data.totalRecords,
    ),
}


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up Sonarr sensors based on a config entry."""
    coordinators: dict[str, SonarrDataUpdateCoordinator[Any]] = hass.data[DOMAIN][
        entry.entry_id
    ]
    async_add_entities(
        SonarrSensor(coordinators[coordinator_type], description)
        for coordinator_type, description in SENSOR_TYPES.items()
    )


class SonarrSensor(SonarrEntity[SonarrDataT], SensorEntity):
    """Implementation of the Sonarr sensor."""

    coordinator: SonarrDataUpdateCoordinator
    entity_description: SonarrSensorEntityDescription[SonarrDataT]

    @property
    def extra_state_attributes(self) -> dict[str, str] | None:
        """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

    @property
    def native_value(self) -> StateType:
        """Return the state of the sensor."""
        return self.entity_description.value_fn(self.coordinator.data)