core/homeassistant/components/sonarr/coordinator.py

"""Data update coordinator for the Sonarr integration."""
from __future__ import annotations

from datetime import timedelta
from typing import TypeVar, Union, cast

from aiopyarr import (
    Command,
    Diskspace,
    SonarrCalendar,
    SonarrQueue,
    SonarrSeries,
    SonarrWantedMissing,
    SystemStatus,
    exceptions,
)
from aiopyarr.models.host_configuration import PyArrHostConfiguration
from aiopyarr.sonarr_client import SonarrClient

from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
import homeassistant.util.dt as dt_util

from .const import CONF_UPCOMING_DAYS, CONF_WANTED_MAX_ITEMS, DOMAIN, LOGGER

SonarrDataT = TypeVar(
    "SonarrDataT",
    bound=Union[
        list[SonarrCalendar],
        list[Command],
        list[Diskspace],
        SonarrQueue,
        list[SonarrSeries],
        SystemStatus,
        SonarrWantedMissing,
    ],
)


class SonarrDataUpdateCoordinator(DataUpdateCoordinator[SonarrDataT]):
    """Data update coordinator for the Sonarr integration."""

    config_entry: ConfigEntry

    def __init__(
        self,
        hass: HomeAssistant,
        host_configuration: PyArrHostConfiguration,
        api_client: SonarrClient,
    ) -> None:
        """Initialize the coordinator."""
        super().__init__(
            hass=hass,
            logger=LOGGER,
            name=DOMAIN,
            update_interval=timedelta(seconds=30),
        )
        self.api_client = api_client
        self.host_configuration = host_configuration
        self.system_version: str | None = None

    async def _async_update_data(self) -> SonarrDataT:
        """Get the latest data from Sonarr."""
        try:
            return await self._fetch_data()

        except exceptions.ArrConnectionException as ex:
            raise UpdateFailed(ex) from ex
        except exceptions.ArrAuthenticationException as ex:
            raise ConfigEntryAuthFailed(
                "API Key is no longer valid. Please reauthenticate"
            ) from ex

    async def _fetch_data(self) -> SonarrDataT:
        """Fetch the actual data."""
        raise NotImplementedError


class CalendarDataUpdateCoordinator(SonarrDataUpdateCoordinator[list[SonarrCalendar]]):
    """Calendar update coordinator."""

    async def _fetch_data(self) -> list[SonarrCalendar]:
        """Fetch the movies data."""
        local = dt_util.start_of_local_day().replace(microsecond=0)
        start = dt_util.as_utc(local)
        end = start + timedelta(days=self.config_entry.options[CONF_UPCOMING_DAYS])
        return cast(
            list[SonarrCalendar],
            await self.api_client.async_get_calendar(
                start_date=start, end_date=end, include_series=True
            ),
        )


class CommandsDataUpdateCoordinator(SonarrDataUpdateCoordinator[list[Command]]):
    """Commands update coordinator for Sonarr."""

    async def _fetch_data(self) -> list[Command]:
        """Fetch the data."""
        return cast(list[Command], await self.api_client.async_get_commands())


class DiskSpaceDataUpdateCoordinator(SonarrDataUpdateCoordinator[list[Diskspace]]):
    """Disk space update coordinator for Sonarr."""

    async def _fetch_data(self) -> list[Diskspace]:
        """Fetch the data."""
        return await self.api_client.async_get_diskspace()


class QueueDataUpdateCoordinator(SonarrDataUpdateCoordinator[SonarrQueue]):
    """Queue update coordinator."""

    async def _fetch_data(self) -> SonarrQueue:
        """Fetch the data."""
        return await self.api_client.async_get_queue(
            include_series=True, include_episode=True
        )


class SeriesDataUpdateCoordinator(SonarrDataUpdateCoordinator[list[SonarrSeries]]):
    """Series update coordinator."""

    async def _fetch_data(self) -> list[SonarrSeries]:
        """Fetch the data."""
        return cast(list[SonarrSeries], await self.api_client.async_get_series())


class StatusDataUpdateCoordinator(SonarrDataUpdateCoordinator[SystemStatus]):
    """Status update coordinator for Sonarr."""

    async def _fetch_data(self) -> SystemStatus:
        """Fetch the data."""
        return await self.api_client.async_get_system_status()


class WantedDataUpdateCoordinator(SonarrDataUpdateCoordinator[SonarrWantedMissing]):
    """Wanted update coordinator."""

    async def _fetch_data(self) -> SonarrWantedMissing:
        """Fetch the data."""
        return await self.api_client.async_get_wanted(
            page_size=self.config_entry.options[CONF_WANTED_MAX_ITEMS],
            include_series=True,
        )