Use aiopyarr for sonarr (#65349)

This commit is contained in:
Chris Talkington 2022-02-22 11:33:10 -06:00 committed by GitHub
parent c14912471d
commit f30681dae7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 464 additions and 345 deletions

View file

@ -7,7 +7,9 @@ from functools import wraps
import logging
from typing import Any, TypeVar
from sonarr import Sonarr, SonarrConnectionError, SonarrError
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
@ -18,7 +20,14 @@ 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_SONARR, DOMAIN
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__)
@ -77,11 +86,22 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up Sonarr sensors based on a config entry."""
sonarr: Sonarr = hass.data[DOMAIN][entry.entry_id][DATA_SONARR]
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, entry.entry_id, description, options)
SonarrSensor(
sonarr,
host_config,
system_status,
entry.entry_id,
description,
options,
)
for description in SENSOR_TYPES
]
@ -102,12 +122,12 @@ def sonarr_exception_handler(
try:
await func(self, *args, **kwargs)
self.last_update_success = True
except SonarrConnectionError as error:
if self.available:
except ArrConnectionException as error:
if self.last_update_success:
_LOGGER.error("Error communicating with API: %s", error)
self.last_update_success = False
except SonarrError as error:
if self.available:
except ArrException as error:
if self.last_update_success:
_LOGGER.error("Invalid response from API: %s", error)
self.last_update_success = False
@ -124,7 +144,9 @@ class SonarrSensor(SonarrEntity, SensorEntity):
def __init__(
self,
sonarr: Sonarr,
sonarr: SonarrClient,
host_config: PyArrHostConfiguration,
system_status: SystemStatus,
entry_id: str,
description: SensorEntityDescription,
options: dict[str, Any],
@ -140,6 +162,8 @@ class SonarrSensor(SonarrEntity, SensorEntity):
super().__init__(
sonarr=sonarr,
host_config=host_config,
system_status=system_status,
entry_id=entry_id,
device_id=entry_id,
)
@ -155,23 +179,30 @@ class SonarrSensor(SonarrEntity, SensorEntity):
key = self.entity_description.key
if key == "diskspace":
await self.sonarr.update()
self.data[key] = await self.sonarr.async_get_diskspace()
elif key == "commands":
self.data[key] = await self.sonarr.commands()
self.data[key] = await self.sonarr.async_get_commands()
elif key == "queue":
self.data[key] = await self.sonarr.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.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.calendar(
start=start.isoformat(), end=end.isoformat()
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.wanted(page_size=self.wanted_max_items)
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:
@ -179,10 +210,10 @@ class SonarrSensor(SonarrEntity, SensorEntity):
attrs = {}
key = self.entity_description.key
if key == "diskspace":
for disk in self.sonarr.app.disks:
free = disk.free / 1024**3
total = disk.total / 1024**3
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[
@ -190,23 +221,33 @@ class SonarrSensor(SonarrEntity, SensorEntity):
] = 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.state
attrs[command.name] = command.status
elif key == "queue" and self.data.get(key) is not None:
for item in self.data[key]:
remaining = 1 if item.size == 0 else item.size_remaining / item.size
for item in self.data[key].records:
remaining = 1 if item.size == 0 else item.sizeleft / item.size
remaining_pct = 100 * (1 - remaining)
name = f"{item.episode.series.title} {item.episode.identifier}"
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]:
attrs[item.series.title] = f"{item.downloaded}/{item.episodes} Episodes"
stats = item.statistics
attrs[
item.title
] = f"{stats.episodeFileCount}/{stats.episodeCount} Episodes"
elif key == "upcoming" and self.data.get(key) is not None:
for episode in self.data[key]:
attrs[episode.series.title] = episode.identifier
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 episode in self.data[key].episodes:
name = f"{episode.series.title} {episode.identifier}"
attrs[name] = episode.airdate
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
@ -215,8 +256,8 @@ class SonarrSensor(SonarrEntity, SensorEntity):
"""Return the state of the sensor."""
key = self.entity_description.key
if key == "diskspace":
total_free = sum(disk.free for disk in self.sonarr.app.disks)
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}"
@ -224,7 +265,7 @@ class SonarrSensor(SonarrEntity, SensorEntity):
return len(self.data[key])
if key == "queue" and self.data.get(key) is not None:
return len(self.data[key])
return self.data[key].totalRecords
if key == "series" and self.data.get(key) is not None:
return len(self.data[key])
@ -233,6 +274,6 @@ class SonarrSensor(SonarrEntity, SensorEntity):
return len(self.data[key])
if key == "wanted" and self.data.get(key) is not None:
return self.data[key].total
return self.data[key].totalRecords
return None