diff --git a/homeassistant/components/rainmachine/__init__.py b/homeassistant/components/rainmachine/__init__.py index 0891d22b641..cfbc95cf009 100644 --- a/homeassistant/components/rainmachine/__init__.py +++ b/homeassistant/components/rainmachine/__init__.py @@ -6,7 +6,7 @@ from collections.abc import Callable, Coroutine from dataclasses import dataclass from datetime import timedelta from functools import partial, wraps -from typing import Any +from typing import Any, cast from regenmaschine import Client from regenmaschine.controller import Controller @@ -169,9 +169,12 @@ COORDINATOR_UPDATE_INTERVAL_MAP = { } +type RainMachineConfigEntry = ConfigEntry[RainMachineData] + + @dataclass class RainMachineData: - """Define an object to be stored in `hass.data`.""" + """Define an object to be stored in `entry.runtime_data`.""" controller: Controller coordinators: dict[str, RainMachineDataUpdateCoordinator] @@ -180,7 +183,7 @@ class RainMachineData: @callback def async_get_entry_for_service_call( hass: HomeAssistant, call: ServiceCall -) -> ConfigEntry: +) -> RainMachineConfigEntry: """Get the controller related to a service call (by device ID).""" device_id = call.data[CONF_DEVICE_ID] device_registry = dr.async_get(hass) @@ -192,20 +195,20 @@ def async_get_entry_for_service_call( if (entry := hass.config_entries.async_get_entry(entry_id)) is None: continue if entry.domain == DOMAIN: - return entry + return cast(RainMachineConfigEntry, entry) raise ValueError(f"No controller for device ID: {device_id}") async def async_update_programs_and_zones( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: RainMachineConfigEntry ) -> None: """Update program and zone DataUpdateCoordinators. Program and zone updates always go together because of how linked they are: programs affect zones and certain combinations of zones affect programs. """ - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data # No gather here to allow http keep-alive to reuse # the connection for each coordinator. await data.coordinators[DATA_PROGRAMS].async_refresh() @@ -213,7 +216,7 @@ async def async_update_programs_and_zones( async def async_setup_entry( # noqa: C901 - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: RainMachineConfigEntry ) -> bool: """Set up RainMachine as config entry.""" websession = aiohttp_client.async_get_clientsession(hass) @@ -315,8 +318,7 @@ async def async_setup_entry( # noqa: C901 # connection for each coordinator. await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {}) - hass.data[DOMAIN][entry.entry_id] = RainMachineData( + entry.runtime_data = RainMachineData( controller=controller, coordinators=coordinators ) @@ -341,7 +343,7 @@ async def async_setup_entry( # noqa: C901 async def wrapper(call: ServiceCall) -> None: """Wrap the service function.""" entry = async_get_entry_for_service_call(hass, call) - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data try: await func(call, data.controller) @@ -461,16 +463,15 @@ async def async_setup_entry( # noqa: C901 return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry( + hass: HomeAssistant, entry: RainMachineConfigEntry +) -> bool: """Unload an RainMachine config entry.""" unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - loaded_entries = [ entry for entry in hass.config_entries.async_entries(DOMAIN) - if entry.state == ConfigEntryState.LOADED + if entry.state is ConfigEntryState.LOADED ] if len(loaded_entries) == 1: # If this is the last loaded instance of RainMachine, deregister any services @@ -489,7 +490,9 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return unload_ok -async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_migrate_entry( + hass: HomeAssistant, entry: RainMachineConfigEntry +) -> bool: """Migrate an old config entry.""" version = entry.version @@ -521,7 +524,9 @@ async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None: +async def async_reload_entry( + hass: HomeAssistant, entry: RainMachineConfigEntry +) -> None: """Handle an options update.""" await hass.config_entries.async_reload(entry.entry_id) @@ -533,7 +538,7 @@ class RainMachineEntity(CoordinatorEntity[RainMachineDataUpdateCoordinator]): def __init__( self, - entry: ConfigEntry, + entry: RainMachineConfigEntry, data: RainMachineData, description: RainMachineEntityDescription, ) -> None: diff --git a/homeassistant/components/rainmachine/binary_sensor.py b/homeassistant/components/rainmachine/binary_sensor.py index 866cddbabbd..574f458ec47 100644 --- a/homeassistant/components/rainmachine/binary_sensor.py +++ b/homeassistant/components/rainmachine/binary_sensor.py @@ -7,13 +7,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, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RainMachineData, RainMachineEntity -from .const import DATA_PROVISION_SETTINGS, DATA_RESTRICTIONS_CURRENT, DOMAIN +from . import RainMachineConfigEntry, RainMachineEntity +from .const import DATA_PROVISION_SETTINGS, DATA_RESTRICTIONS_CURRENT from .model import RainMachineEntityDescription from .util import ( EntityDomainReplacementStrategy, @@ -93,10 +92,12 @@ BINARY_SENSOR_DESCRIPTIONS = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: RainMachineConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up RainMachine binary sensors based on a config entry.""" - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data async_finish_entity_domain_replacements( hass, @@ -125,15 +126,13 @@ async def async_setup_entry( } async_add_entities( - [ - api_category_sensor_map[description.api_category](entry, data, description) - for description in BINARY_SENSOR_DESCRIPTIONS - if ( - (coordinator := data.coordinators[description.api_category]) is not None - and coordinator.data - and key_exists(coordinator.data, description.data_key) - ) - ] + api_category_sensor_map[description.api_category](entry, data, description) + for description in BINARY_SENSOR_DESCRIPTIONS + if ( + (coordinator := data.coordinators[description.api_category]) is not None + and coordinator.data + and key_exists(coordinator.data, description.data_key) + ) ) diff --git a/homeassistant/components/rainmachine/button.py b/homeassistant/components/rainmachine/button.py index 24486a34b88..7087e5e5b8e 100644 --- a/homeassistant/components/rainmachine/button.py +++ b/homeassistant/components/rainmachine/button.py @@ -13,15 +13,14 @@ from homeassistant.components.button import ( ButtonEntity, ButtonEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RainMachineData, RainMachineEntity -from .const import DATA_PROVISION_SETTINGS, DOMAIN +from . import RainMachineConfigEntry, RainMachineEntity +from .const import DATA_PROVISION_SETTINGS from .model import RainMachineEntityDescription @@ -52,11 +51,12 @@ BUTTON_DESCRIPTIONS = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: RainMachineConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up RainMachine buttons based on a config entry.""" - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] - + data = entry.runtime_data async_add_entities( RainMachineButton(entry, data, description) for description in BUTTON_DESCRIPTIONS diff --git a/homeassistant/components/rainmachine/coordinator.py b/homeassistant/components/rainmachine/coordinator.py index 620bdb2da9b..df7972ef31d 100644 --- a/homeassistant/components/rainmachine/coordinator.py +++ b/homeassistant/components/rainmachine/coordinator.py @@ -4,9 +4,8 @@ from __future__ import annotations from collections.abc import Callable, Coroutine from datetime import timedelta -from typing import Any +from typing import TYPE_CHECKING, Any -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import ( async_dispatcher_connect, @@ -19,17 +18,20 @@ from .const import LOGGER SIGNAL_REBOOT_COMPLETED = "rainmachine_reboot_completed_{0}" SIGNAL_REBOOT_REQUESTED = "rainmachine_reboot_requested_{0}" +if TYPE_CHECKING: + from . import RainMachineConfigEntry + class RainMachineDataUpdateCoordinator(DataUpdateCoordinator[dict]): """Define an extended DataUpdateCoordinator.""" - config_entry: ConfigEntry + config_entry: RainMachineConfigEntry def __init__( self, hass: HomeAssistant, *, - entry: ConfigEntry, + entry: RainMachineConfigEntry, name: str, api_category: str, update_interval: timedelta, diff --git a/homeassistant/components/rainmachine/diagnostics.py b/homeassistant/components/rainmachine/diagnostics.py index 5564ee693a4..598b8aefa5f 100644 --- a/homeassistant/components/rainmachine/diagnostics.py +++ b/homeassistant/components/rainmachine/diagnostics.py @@ -7,7 +7,6 @@ from typing import Any from regenmaschine.errors import RainMachineError from homeassistant.components.diagnostics import async_redact_data -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONF_ELEVATION, CONF_LATITUDE, @@ -17,8 +16,8 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant -from . import RainMachineData -from .const import DOMAIN, LOGGER +from . import RainMachineConfigEntry +from .const import LOGGER CONF_STATION_ID = "stationID" CONF_STATION_NAME = "stationName" @@ -40,10 +39,10 @@ TO_REDACT = { async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: RainMachineConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data try: controller_diagnostics = await data.controller.diagnostics.current() diff --git a/homeassistant/components/rainmachine/select.py b/homeassistant/components/rainmachine/select.py index bb622330897..73de33cc8ed 100644 --- a/homeassistant/components/rainmachine/select.py +++ b/homeassistant/components/rainmachine/select.py @@ -14,8 +14,8 @@ from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM, UnitSystem -from . import RainMachineData, RainMachineEntity -from .const import DATA_RESTRICTIONS_UNIVERSAL, DOMAIN +from . import RainMachineConfigEntry, RainMachineData, RainMachineEntity +from .const import DATA_RESTRICTIONS_UNIVERSAL from .model import RainMachineEntityDescription from .util import key_exists @@ -81,10 +81,12 @@ SELECT_DESCRIPTIONS = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: RainMachineConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up RainMachine selects based on a config entry.""" - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data entity_map = { TYPE_FREEZE_PROTECTION_TEMPERATURE: FreezeProtectionTemperatureSelect, diff --git a/homeassistant/components/rainmachine/sensor.py b/homeassistant/components/rainmachine/sensor.py index 15188e86963..5363000a8ac 100644 --- a/homeassistant/components/rainmachine/sensor.py +++ b/homeassistant/components/rainmachine/sensor.py @@ -20,8 +20,8 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.util.dt import utc_from_timestamp, utcnow -from . import RainMachineData, RainMachineEntity -from .const import DATA_PROGRAMS, DATA_PROVISION_SETTINGS, DATA_ZONES, DOMAIN +from . import RainMachineConfigEntry, RainMachineData, RainMachineEntity +from .const import DATA_PROGRAMS, DATA_PROVISION_SETTINGS, DATA_ZONES from .model import RainMachineEntityDescription from .util import ( RUN_STATE_MAP, @@ -151,10 +151,12 @@ SENSOR_DESCRIPTIONS = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: RainMachineConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up RainMachine sensors based on a config entry.""" - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data async_finish_entity_domain_replacements( hass, diff --git a/homeassistant/components/rainmachine/switch.py b/homeassistant/components/rainmachine/switch.py index 667e609e11c..d4c0064219e 100644 --- a/homeassistant/components/rainmachine/switch.py +++ b/homeassistant/components/rainmachine/switch.py @@ -20,7 +20,12 @@ from homeassistant.helpers import config_validation as cv, entity_platform from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import VolDictType -from . import RainMachineData, RainMachineEntity, async_update_programs_and_zones +from . import ( + RainMachineConfigEntry, + RainMachineData, + RainMachineEntity, + async_update_programs_and_zones, +) from .const import ( CONF_ALLOW_INACTIVE_ZONES_TO_RUN, CONF_DEFAULT_ZONE_RUN_TIME, @@ -31,7 +36,6 @@ from .const import ( DATA_RESTRICTIONS_UNIVERSAL, DATA_ZONES, DEFAULT_ZONE_RUN, - DOMAIN, ) from .model import RainMachineEntityDescription from .util import RUN_STATE_MAP, key_exists @@ -173,7 +177,9 @@ RESTRICTIONS_SWITCH_DESCRIPTIONS = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: RainMachineConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up RainMachine switches based on a config entry.""" platform = entity_platform.async_get_current_platform() @@ -195,7 +201,7 @@ async def async_setup_entry( schema_dict = cast(VolDictType, schema) platform.async_register_entity_service(service_name, schema_dict, method) - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data entities: list[RainMachineBaseSwitch] = [] for kind, api_category, switch_class, switch_enabled_class in ( diff --git a/homeassistant/components/rainmachine/update.py b/homeassistant/components/rainmachine/update.py index 38bf74effa0..a7c11061718 100644 --- a/homeassistant/components/rainmachine/update.py +++ b/homeassistant/components/rainmachine/update.py @@ -12,13 +12,12 @@ from homeassistant.components.update import ( UpdateEntity, UpdateEntityFeature, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RainMachineData, RainMachineEntity -from .const import DATA_MACHINE_FIRMWARE_UPDATE_STATUS, DOMAIN +from . import RainMachineConfigEntry, RainMachineEntity +from .const import DATA_MACHINE_FIRMWARE_UPDATE_STATUS from .model import RainMachineEntityDescription @@ -50,11 +49,12 @@ UPDATE_DESCRIPTION = RainMachineEntityDescription( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: RainMachineConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up Rainmachine update based on a config entry.""" - data: RainMachineData = hass.data[DOMAIN][entry.entry_id] - + data = entry.runtime_data async_add_entities([RainMachineUpdateEntity(entry, data, UPDATE_DESCRIPTION)])