From dcca749d50d047e3dc0c66c7793ecd73dd03752f Mon Sep 17 00:00:00 2001 From: Robert Hillis Date: Mon, 17 Jun 2024 07:47:49 -0400 Subject: [PATCH] Store runtime data inside the config entry in Radarr (#119749) * Store runtime data inside the config entry in Radarr * move entry typing outside constructor --- homeassistant/components/radarr/__init__.py | 46 ++++++++++++------- .../components/radarr/binary_sensor.py | 9 ++-- homeassistant/components/radarr/calendar.py | 10 ++-- .../components/radarr/config_flow.py | 8 ++-- .../components/radarr/coordinator.py | 8 ++-- homeassistant/components/radarr/sensor.py | 11 ++--- 6 files changed, 49 insertions(+), 43 deletions(-) diff --git a/homeassistant/components/radarr/__init__.py b/homeassistant/components/radarr/__init__.py index d3e44e6b7fc..b528e701c71 100644 --- a/homeassistant/components/radarr/__init__.py +++ b/homeassistant/components/radarr/__init__.py @@ -2,7 +2,8 @@ from __future__ import annotations -from typing import Any, cast +from dataclasses import dataclass, fields +from typing import cast from aiopyarr.models.host_configuration import PyArrHostConfiguration from aiopyarr.radarr_client import RadarrClient @@ -34,9 +35,22 @@ from .coordinator import ( ) PLATFORMS = [Platform.BINARY_SENSOR, Platform.CALENDAR, Platform.SENSOR] +type RadarrConfigEntry = ConfigEntry[RadarrData] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +@dataclass(kw_only=True, slots=True) +class RadarrData: + """Radarr data type.""" + + calendar: CalendarUpdateCoordinator + disk_space: DiskSpaceDataUpdateCoordinator + health: HealthDataUpdateCoordinator + movie: MoviesDataUpdateCoordinator + queue: QueueDataUpdateCoordinator + status: StatusDataUpdateCoordinator + + +async def async_setup_entry(hass: HomeAssistant, entry: RadarrConfigEntry) -> bool: """Set up Radarr from a config entry.""" host_configuration = PyArrHostConfiguration( api_token=entry.data[CONF_API_KEY], @@ -47,27 +61,25 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: host_configuration=host_configuration, session=async_get_clientsession(hass, entry.data[CONF_VERIFY_SSL]), ) - coordinators: dict[str, RadarrDataUpdateCoordinator[Any]] = { - "calendar": CalendarUpdateCoordinator(hass, host_configuration, radarr), - "disk_space": DiskSpaceDataUpdateCoordinator(hass, host_configuration, radarr), - "health": HealthDataUpdateCoordinator(hass, host_configuration, radarr), - "movie": MoviesDataUpdateCoordinator(hass, host_configuration, radarr), - "queue": QueueDataUpdateCoordinator(hass, host_configuration, radarr), - "status": StatusDataUpdateCoordinator(hass, host_configuration, radarr), - } - for coordinator in coordinators.values(): - await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinators + data = RadarrData( + calendar=CalendarUpdateCoordinator(hass, host_configuration, radarr), + disk_space=DiskSpaceDataUpdateCoordinator(hass, host_configuration, radarr), + health=HealthDataUpdateCoordinator(hass, host_configuration, radarr), + movie=MoviesDataUpdateCoordinator(hass, host_configuration, radarr), + queue=QueueDataUpdateCoordinator(hass, host_configuration, radarr), + status=StatusDataUpdateCoordinator(hass, host_configuration, radarr), + ) + for field in fields(data): + await getattr(data, field.name).async_config_entry_first_refresh() + entry.runtime_data = data await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: RadarrConfigEntry) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - hass.data[DOMAIN].pop(entry.entry_id) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) class RadarrEntity(CoordinatorEntity[RadarrDataUpdateCoordinator[T]]): diff --git a/homeassistant/components/radarr/binary_sensor.py b/homeassistant/components/radarr/binary_sensor.py index 4962ef81614..6c0468cff58 100644 --- a/homeassistant/components/radarr/binary_sensor.py +++ b/homeassistant/components/radarr/binary_sensor.py @@ -9,13 +9,12 @@ from homeassistant.components.binary_sensor import ( BinarySensorEntity, BinarySensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RadarrEntity -from .const import DOMAIN, HEALTH_ISSUES +from . import RadarrConfigEntry, RadarrEntity +from .const import HEALTH_ISSUES BINARY_SENSOR_TYPE = BinarySensorEntityDescription( key="health", @@ -27,11 +26,11 @@ BINARY_SENSOR_TYPE = BinarySensorEntityDescription( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: RadarrConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Radarr sensors based on a config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id]["health"] + coordinator = entry.runtime_data.health async_add_entities([RadarrBinarySensor(coordinator, BINARY_SENSOR_TYPE)]) diff --git a/homeassistant/components/radarr/calendar.py b/homeassistant/components/radarr/calendar.py index ad5e1b8ffd9..4f866123a1a 100644 --- a/homeassistant/components/radarr/calendar.py +++ b/homeassistant/components/radarr/calendar.py @@ -5,13 +5,11 @@ from __future__ import annotations from datetime import datetime from homeassistant.components.calendar import CalendarEntity, CalendarEvent -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity import EntityDescription from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RadarrEntity -from .const import DOMAIN +from . import RadarrConfigEntry, RadarrEntity from .coordinator import CalendarUpdateCoordinator, RadarrEvent CALENDAR_TYPE = EntityDescription( @@ -21,10 +19,12 @@ CALENDAR_TYPE = EntityDescription( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: RadarrConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Radarr calendar entity.""" - coordinator = hass.data[DOMAIN][entry.entry_id]["calendar"] + coordinator = entry.runtime_data.calendar async_add_entities([RadarrCalendarEntity(coordinator, CALENDAR_TYPE)]) diff --git a/homeassistant/components/radarr/config_flow.py b/homeassistant/components/radarr/config_flow.py index 81589c5fe30..3bf0796a9a8 100644 --- a/homeassistant/components/radarr/config_flow.py +++ b/homeassistant/components/radarr/config_flow.py @@ -11,11 +11,12 @@ from aiopyarr.models.host_configuration import PyArrHostConfiguration from aiopyarr.radarr_client import RadarrClient import voluptuous as vol -from homeassistant.config_entries import ConfigEntry, ConfigFlow, ConfigFlowResult +from homeassistant.config_entries import ConfigFlow, ConfigFlowResult from homeassistant.const import CONF_API_KEY, CONF_URL, CONF_VERIFY_SSL from homeassistant.core import HomeAssistant from homeassistant.helpers.aiohttp_client import async_get_clientsession +from . import RadarrConfigEntry from .const import DEFAULT_NAME, DEFAULT_URL, DOMAIN @@ -23,10 +24,7 @@ class RadarrConfigFlow(ConfigFlow, domain=DOMAIN): """Handle a config flow for Radarr.""" VERSION = 1 - - def __init__(self) -> None: - """Initialize the flow.""" - self.entry: ConfigEntry | None = None + entry: RadarrConfigEntry | None = None async def async_step_reauth(self, _: Mapping[str, Any]) -> ConfigFlowResult: """Handle configuration by re-auth.""" diff --git a/homeassistant/components/radarr/coordinator.py b/homeassistant/components/radarr/coordinator.py index 47a1862b8ae..6e8a3d55d3e 100644 --- a/homeassistant/components/radarr/coordinator.py +++ b/homeassistant/components/radarr/coordinator.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod import asyncio from dataclasses import dataclass from datetime import date, datetime, timedelta -from typing import Generic, TypeVar, cast +from typing import TYPE_CHECKING, Generic, TypeVar, cast from aiopyarr import ( Health, @@ -20,13 +20,15 @@ from aiopyarr.models.host_configuration import PyArrHostConfiguration from aiopyarr.radarr_client import RadarrClient from homeassistant.components.calendar import CalendarEvent -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import DEFAULT_MAX_RECORDS, DOMAIN, LOGGER +if TYPE_CHECKING: + from . import RadarrConfigEntry + T = TypeVar("T", bound=SystemStatus | list[RootFolder] | list[Health] | int | None) @@ -45,7 +47,7 @@ class RadarrEvent(CalendarEvent, RadarrEventMixIn): class RadarrDataUpdateCoordinator(DataUpdateCoordinator[T], Generic[T], ABC): """Data update coordinator for the Radarr integration.""" - config_entry: ConfigEntry + config_entry: RadarrConfigEntry _update_interval = timedelta(seconds=30) def __init__( diff --git a/homeassistant/components/radarr/sensor.py b/homeassistant/components/radarr/sensor.py index e6700fb3637..441c44de781 100644 --- a/homeassistant/components/radarr/sensor.py +++ b/homeassistant/components/radarr/sensor.py @@ -15,13 +15,11 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory, UnitOfInformation from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RadarrEntity -from .const import DOMAIN +from . import RadarrConfigEntry, RadarrEntity from .coordinator import RadarrDataUpdateCoordinator, T @@ -117,16 +115,13 @@ PARALLEL_UPDATES = 1 async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: RadarrConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Radarr sensors based on a config entry.""" - coordinators: dict[str, RadarrDataUpdateCoordinator[Any]] = hass.data[DOMAIN][ - entry.entry_id - ] entities: list[RadarrSensor[Any]] = [] for coordinator_type, description in SENSOR_TYPES.items(): - coordinator = coordinators[coordinator_type] + coordinator = getattr(entry.runtime_data, coordinator_type) if coordinator_type != "disk_space": entities.append(RadarrSensor(coordinator, description)) else: