Convert rainmachine to use entry.runtime_data (#122532)

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:

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)
+        )
     )
 
 

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

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,

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()

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,

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,

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 (

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)])