"""Support for Sonarr sensors."""
from __future__ import annotations
from collections.abc import Awaitable, Callable, Coroutine
from datetime import timedelta
from functools import wraps
import logging
from typing import Any, TypeVar
from aiopyarr import ArrConnectionException, ArrException, SystemStatus
from aiopyarr.models.host_configuration import PyArrHostConfiguration
from aiopyarr.sonarr_client import SonarrClient
from typing_extensions import Concatenate, ParamSpec
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_HOST_CONFIG,
DATA_SONARR,
DATA_SYSTEM_STATUS,
DOMAIN,
)
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,
),
)
_SonarrSensorT = TypeVar("_SonarrSensorT", bound="SonarrSensor")
_P = ParamSpec("_P")
async def async_setup_entry(
hass: HomeAssistant,
entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up Sonarr sensors based on a config entry."""
sonarr: SonarrClient = hass.data[DOMAIN][entry.entry_id][DATA_SONARR]
host_config: PyArrHostConfiguration = hass.data[DOMAIN][entry.entry_id][
DATA_HOST_CONFIG
]
system_status: SystemStatus = hass.data[DOMAIN][entry.entry_id][DATA_SYSTEM_STATUS]
options: dict[str, Any] = dict(entry.options)
entities = [
SonarrSensor(
sonarr,
host_config,
system_status,
entry.entry_id,
description,
options,
)
for description in SENSOR_TYPES
]
async_add_entities(entities, True)
def sonarr_exception_handler(
func: Callable[Concatenate[_SonarrSensorT, _P], Awaitable[None]]
) -> Callable[Concatenate[_SonarrSensorT, _P], Coroutine[Any, Any, None]]:
"""Decorate Sonarr calls to handle Sonarr exceptions.
A decorator that wraps the passed in function, catches Sonarr errors,
and handles the availability of the entity.
"""
@wraps(func)
async def wrapper(
self: _SonarrSensorT, *args: _P.args, **kwargs: _P.kwargs
) -> None:
try:
await func(self, *args, **kwargs)
self.last_update_success = True
except ArrConnectionException as error:
if self.last_update_success:
_LOGGER.error("Error communicating with API: %s", error)
self.last_update_success = False
except ArrException as error:
if self.last_update_success:
_LOGGER.error("Invalid response from API: %s", error)
self.last_update_success = False
return wrapper
class SonarrSensor(SonarrEntity, SensorEntity):
"""Implementation of the Sonarr sensor."""
data: dict[str, Any]
last_update_success: bool
upcoming_days: int
wanted_max_items: int
def __init__(
self,
sonarr: SonarrClient,
host_config: PyArrHostConfiguration,
system_status: SystemStatus,
entry_id: str,
description: SensorEntityDescription,
options: dict[str, Any],
) -> None:
"""Initialize Sonarr sensor."""
self.entity_description = description
self._attr_unique_id = f"{entry_id}_{description.key}"
self.data = {}
self.last_update_success = True
self.upcoming_days = options[CONF_UPCOMING_DAYS]
self.wanted_max_items = options[CONF_WANTED_MAX_ITEMS]
super().__init__(
sonarr=sonarr,
host_config=host_config,
system_status=system_status,
entry_id=entry_id,
device_id=entry_id,
)
@property
def available(self) -> bool:
"""Return sensor availability."""
return self.last_update_success
@sonarr_exception_handler
async def async_update(self) -> None:
"""Update entity."""
key = self.entity_description.key
if key == "diskspace":
self.data[key] = await self.sonarr.async_get_diskspace()
elif key == "commands":
self.data[key] = await self.sonarr.async_get_commands()
elif key == "queue":
self.data[key] = await self.sonarr.async_get_queue(
include_series=True, include_episode=True
)
elif key == "series":
self.data[key] = await self.sonarr.async_get_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.async_get_calendar(
start_date=start,
end_date=end,
include_series=True,
)
elif key == "wanted":
self.data[key] = await self.sonarr.async_get_wanted(
page_size=self.wanted_max_items,
include_series=True,
)
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the entity."""
attrs = {}
key = self.entity_description.key
if key == "diskspace" and self.data.get(key) is not None:
for disk in self.data[key]:
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" and self.data.get(key) is not None:
for command in self.data[key]:
attrs[command.name] = command.status
elif key == "queue" and self.data.get(key) is not None:
for item in self.data[key].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" and self.data.get(key) is not None:
for item in self.data[key]:
stats = item.statistics
attrs[
item.title
] = f"{getattr(stats,'episodeFileCount', 0)}/{getattr(stats, 'episodeCount', 0)} Episodes"
elif key == "upcoming" and self.data.get(key) is not None:
for episode in self.data[key]:
identifier = f"S{episode.seasonNumber:02d}E{episode.episodeNumber:02d}"
attrs[episode.series.title] = identifier
elif key == "wanted" and self.data.get(key) is not None:
for item in self.data[key].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."""
key = self.entity_description.key
if key == "diskspace" and self.data.get(key) is not None:
total_free = sum(disk.freeSpace for disk in self.data[key])
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 self.data[key].totalRecords
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].totalRecords
return None