mirror of
https://github.com/home-assistant/core
synced 2024-08-02 23:40:32 +02:00
282 lines
9.5 KiB
Python
282 lines
9.5 KiB
Python
"""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
|