core/homeassistant/components/sonarr/sensor.py

393 lines
12 KiB
Python

"""Support for Sonarr sensors."""
from datetime import timedelta
import logging
from typing import Any, Callable, Dict, List, Optional
from sonarr import Sonarr, SonarrConnectionError, SonarrError
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import DATA_GIGABYTES
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.typing import HomeAssistantType
import homeassistant.util.dt as dt_util
from . import SonarrEntity
from .const import CONF_UPCOMING_DAYS, CONF_WANTED_MAX_ITEMS, DATA_SONARR, DOMAIN
_LOGGER = logging.getLogger(__name__)
async def async_setup_entry(
hass: HomeAssistantType,
entry: ConfigEntry,
async_add_entities: Callable[[List[Entity], bool], None],
) -> 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):
"""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: Optional[str] = None,
) -> None:
"""Initialize Sonarr sensor."""
self._unit_of_measurement = unit_of_measurement
self._key = key
self._unique_id = f"{entry_id}_{key}"
self.last_update_success = False
super().__init__(
sonarr=sonarr,
entry_id=entry_id,
device_id=entry_id,
name=name,
icon=icon,
enabled_default=enabled_default,
)
@property
def unique_id(self) -> str:
"""Return the unique ID for this sensor."""
return self._unique_id
@property
def available(self) -> bool:
"""Return sensor availability."""
return self.last_update_success
@property
def unit_of_measurement(self) -> str:
"""Return the unit this state is expressed in."""
return self._unit_of_measurement
class SonarrCommandsSensor(SonarrSensor):
"""Defines a Sonarr Commands sensor."""
def __init__(self, sonarr: Sonarr, entry_id: str) -> None:
"""Initialize Sonarr Commands sensor."""
self._commands = []
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 device_state_attributes(self) -> Optional[Dict[str, Any]]:
"""Return the state attributes of the entity."""
attrs = {}
for command in self._commands:
attrs[command.name] = command.state
return attrs
@property
def state(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 = []
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 device_state_attributes(self) -> Optional[Dict[str, Any]]:
"""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 state(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 = []
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 device_state_attributes(self) -> Optional[Dict[str, Any]]:
"""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 state(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 = []
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 device_state_attributes(self) -> Optional[Dict[str, Any]]:
"""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 state(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 = []
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",
)
async def async_added_to_hass(self):
"""Listen for signals."""
await super().async_added_to_hass()
self.async_on_remove(
async_dispatcher_connect(
self.hass,
f"sonarr.{self._entry_id}.entry_options_update",
self.async_update_entry_options,
)
)
@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()
)
async def async_update_entry_options(self, options: dict) -> None:
"""Update sensor settings when config entry options are update."""
self._days = options[CONF_UPCOMING_DAYS]
@property
def device_state_attributes(self) -> Optional[Dict[str, Any]]:
"""Return the state attributes of the entity."""
attrs = {}
for episode in self._upcoming:
attrs[episode.series.title] = episode.identifier
return attrs
@property
def state(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 = None
self._total: Optional[int] = 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,
)
async def async_added_to_hass(self):
"""Listen for signals."""
await super().async_added_to_hass()
self.async_on_remove(
async_dispatcher_connect(
self.hass,
f"sonarr.{self._entry_id}.entry_options_update",
self.async_update_entry_options,
)
)
@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
async def async_update_entry_options(self, options: dict) -> None:
"""Update sensor settings when config entry options are update."""
self._max_items = options[CONF_WANTED_MAX_ITEMS]
@property
def device_state_attributes(self) -> Optional[Dict[str, Any]]:
"""Return the state attributes of the entity."""
attrs = {}
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 state(self) -> Optional[int]:
"""Return the state of the sensor."""
return self._total