core/homeassistant/components/sonarr/sensor.py

360 lines
11 KiB
Python

"""Support for Sonarr sensors."""
from __future__ import annotations
from datetime import timedelta
import logging
from sonarr import Sonarr, SonarrConnectionError, SonarrError
from sonarr.models import (
CommandItem,
Disk,
Episode,
QueueItem,
SeriesItem,
WantedResults,
)
from homeassistant.components.sensor import SensorEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import DATA_GIGABYTES
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.util.dt as dt_util
from .const import CONF_UPCOMING_DAYS, CONF_WANTED_MAX_ITEMS, DATA_SONARR, DOMAIN
from .entity import SonarrEntity
_LOGGER = logging.getLogger(__name__)
async def async_setup_entry(
hass: HomeAssistant,
entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up Sonarr sensors based on a config entry."""
options = entry.options
sonarr = hass.data[DOMAIN][entry.entry_id][DATA_SONARR]
entities = [
SonarrCommandsSensor(sonarr, entry.entry_id),
SonarrDiskspaceSensor(sonarr, entry.entry_id),
SonarrQueueSensor(sonarr, entry.entry_id),
SonarrSeriesSensor(sonarr, entry.entry_id),
SonarrUpcomingSensor(sonarr, entry.entry_id, days=options[CONF_UPCOMING_DAYS]),
SonarrWantedSensor(
sonarr, entry.entry_id, max_items=options[CONF_WANTED_MAX_ITEMS]
),
]
async_add_entities(entities, True)
def sonarr_exception_handler(func):
"""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.
"""
async def handler(self, *args, **kwargs):
try:
await func(self, *args, **kwargs)
self.last_update_success = True
except SonarrConnectionError as error:
if self.available:
_LOGGER.error("Error communicating with API: %s", error)
self.last_update_success = False
except SonarrError as error:
if self.available:
_LOGGER.error("Invalid response from API: %s", error)
self.last_update_success = False
return handler
class SonarrSensor(SonarrEntity, SensorEntity):
"""Implementation of the Sonarr sensor."""
def __init__(
self,
*,
sonarr: Sonarr,
entry_id: str,
enabled_default: bool = True,
icon: str,
key: str,
name: str,
unit_of_measurement: str | None = None,
) -> None:
"""Initialize Sonarr sensor."""
self._key = key
self._attr_name = name
self._attr_icon = icon
self._attr_unique_id = f"{entry_id}_{key}"
self._attr_native_unit_of_measurement = unit_of_measurement
self._attr_entity_registry_enabled_default = enabled_default
self.last_update_success = False
super().__init__(
sonarr=sonarr,
entry_id=entry_id,
device_id=entry_id,
)
@property
def available(self) -> bool:
"""Return sensor availability."""
return self.last_update_success
class SonarrCommandsSensor(SonarrSensor):
"""Defines a Sonarr Commands sensor."""
def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
"""Initialize Sonarr Commands sensor."""
self._commands: list[CommandItem] = []
super().__init__(
sonarr=sonarr,
entry_id=entry_id,
icon="mdi:code-braces",
key="commands",
name=f"{sonarr.app.info.app_name} Commands",
unit_of_measurement="Commands",
enabled_default=False,
)
@sonarr_exception_handler
async def async_update(self) -> None:
"""Update entity."""
self._commands = await self.sonarr.commands()
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the entity."""
attrs = {}
for command in self._commands:
attrs[command.name] = command.state
return attrs
@property
def native_value(self) -> int:
"""Return the state of the sensor."""
return len(self._commands)
class SonarrDiskspaceSensor(SonarrSensor):
"""Defines a Sonarr Disk Space sensor."""
def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
"""Initialize Sonarr Disk Space sensor."""
self._disks: list[Disk] = []
self._total_free = 0
super().__init__(
sonarr=sonarr,
entry_id=entry_id,
icon="mdi:harddisk",
key="diskspace",
name=f"{sonarr.app.info.app_name} Disk Space",
unit_of_measurement=DATA_GIGABYTES,
enabled_default=False,
)
@sonarr_exception_handler
async def async_update(self) -> None:
"""Update entity."""
app = await self.sonarr.update()
self._disks = app.disks
self._total_free = sum(disk.free for disk in self._disks)
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the entity."""
attrs = {}
for disk in self._disks:
free = disk.free / 1024 ** 3
total = disk.total / 1024 ** 3
usage = free / total * 100
attrs[
disk.path
] = f"{free:.2f}/{total:.2f}{self.unit_of_measurement} ({usage:.2f}%)"
return attrs
@property
def native_value(self) -> str:
"""Return the state of the sensor."""
free = self._total_free / 1024 ** 3
return f"{free:.2f}"
class SonarrQueueSensor(SonarrSensor):
"""Defines a Sonarr Queue sensor."""
def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
"""Initialize Sonarr Queue sensor."""
self._queue: list[QueueItem] = []
super().__init__(
sonarr=sonarr,
entry_id=entry_id,
icon="mdi:download",
key="queue",
name=f"{sonarr.app.info.app_name} Queue",
unit_of_measurement="Episodes",
enabled_default=False,
)
@sonarr_exception_handler
async def async_update(self) -> None:
"""Update entity."""
self._queue = await self.sonarr.queue()
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the entity."""
attrs = {}
for item in self._queue:
remaining = 1 if item.size == 0 else item.size_remaining / item.size
remaining_pct = 100 * (1 - remaining)
name = f"{item.episode.series.title} {item.episode.identifier}"
attrs[name] = f"{remaining_pct:.2f}%"
return attrs
@property
def native_value(self) -> int:
"""Return the state of the sensor."""
return len(self._queue)
class SonarrSeriesSensor(SonarrSensor):
"""Defines a Sonarr Series sensor."""
def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
"""Initialize Sonarr Series sensor."""
self._items: list[SeriesItem] = []
super().__init__(
sonarr=sonarr,
entry_id=entry_id,
icon="mdi:television",
key="series",
name=f"{sonarr.app.info.app_name} Shows",
unit_of_measurement="Series",
enabled_default=False,
)
@sonarr_exception_handler
async def async_update(self) -> None:
"""Update entity."""
self._items = await self.sonarr.series()
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the entity."""
attrs = {}
for item in self._items:
attrs[item.series.title] = f"{item.downloaded}/{item.episodes} Episodes"
return attrs
@property
def native_value(self) -> int:
"""Return the state of the sensor."""
return len(self._items)
class SonarrUpcomingSensor(SonarrSensor):
"""Defines a Sonarr Upcoming sensor."""
def __init__(self, sonarr: Sonarr, entry_id: str, days: int = 1) -> None:
"""Initialize Sonarr Upcoming sensor."""
self._days = days
self._upcoming: list[Episode] = []
super().__init__(
sonarr=sonarr,
entry_id=entry_id,
icon="mdi:television",
key="upcoming",
name=f"{sonarr.app.info.app_name} Upcoming",
unit_of_measurement="Episodes",
)
@sonarr_exception_handler
async def async_update(self) -> None:
"""Update entity."""
local = dt_util.start_of_local_day().replace(microsecond=0)
start = dt_util.as_utc(local)
end = start + timedelta(days=self._days)
self._upcoming = await self.sonarr.calendar(
start=start.isoformat(), end=end.isoformat()
)
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the entity."""
attrs = {}
for episode in self._upcoming:
attrs[episode.series.title] = episode.identifier
return attrs
@property
def native_value(self) -> int:
"""Return the state of the sensor."""
return len(self._upcoming)
class SonarrWantedSensor(SonarrSensor):
"""Defines a Sonarr Wanted sensor."""
def __init__(self, sonarr: Sonarr, entry_id: str, max_items: int = 10) -> None:
"""Initialize Sonarr Wanted sensor."""
self._max_items = max_items
self._results: WantedResults | None = None
self._total: int | None = None
super().__init__(
sonarr=sonarr,
entry_id=entry_id,
icon="mdi:television",
key="wanted",
name=f"{sonarr.app.info.app_name} Wanted",
unit_of_measurement="Episodes",
enabled_default=False,
)
@sonarr_exception_handler
async def async_update(self) -> None:
"""Update entity."""
self._results = await self.sonarr.wanted(page_size=self._max_items)
self._total = self._results.total
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the entity."""
attrs: dict[str, str] = {}
if self._results is not None:
for episode in self._results.episodes:
name = f"{episode.series.title} {episode.identifier}"
attrs[name] = episode.airdate
return attrs
@property
def native_value(self) -> int | None:
"""Return the state of the sensor."""
return self._total