Use runtime_data in control4 (#136403)

homeassistant/components/control4/__init__.py

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 import json
 import logging
+from typing import Any
 
 from aiohttp import client_exceptions
 from pyControl4.account import C4Account
@@ -25,14 +27,7 @@
 
 from .const import (
     API_RETRY_TIMES,
-    CONF_ACCOUNT,
-    CONF_CONFIG_LISTENER,
     CONF_CONTROLLER_UNIQUE_ID,
-    CONF_DIRECTOR,
-    CONF_DIRECTOR_ALL_ITEMS,
-    CONF_DIRECTOR_MODEL,
-    CONF_DIRECTOR_SW_VERSION,
-    CONF_UI_CONFIGURATION,
     DEFAULT_SCAN_INTERVAL,
     DOMAIN,
 )
@@ -42,6 +37,23 @@
 PLATFORMS = [Platform.LIGHT, Platform.MEDIA_PLAYER]
 
 
+@dataclass
+class Control4RuntimeData:
+    """Control4 runtime data."""
+
+    account: C4Account
+    controller_unique_id: str
+    director: C4Director
+    director_all_items: list[dict[str, Any]]
+    director_model: str
+    director_sw_version: str
+    scan_interval: int
+    ui_configuration: dict[str, Any] | None
+
+
+type Control4ConfigEntry = ConfigEntry[Control4RuntimeData]
+
+
 async def call_c4_api_retry(func, *func_args):
     """Call C4 API function and retry on failure."""
     # Ruff doesn't understand this loop - the exception is always raised after the retries
@@ -54,10 +66,8 @@ async def call_c4_api_retry(func, *func_args):
                 raise ConfigEntryNotReady(exception) from exception
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: Control4ConfigEntry) -> bool:
     """Set up Control4 from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
-    entry_data = hass.data[DOMAIN].setdefault(entry.entry_id, {})
     account_session = aiohttp_client.async_get_clientsession(hass)
 
     config = entry.data
@@ -76,10 +86,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             exception,
         )
         return False
-    entry_data[CONF_ACCOUNT] = account
 
-    controller_unique_id = config[CONF_CONTROLLER_UNIQUE_ID]
-    entry_data[CONF_CONTROLLER_UNIQUE_ID] = controller_unique_id
+    controller_unique_id: str = config[CONF_CONTROLLER_UNIQUE_ID]
 
     director_token_dict = await call_c4_api_retry(
         account.getDirectorBearerToken, controller_unique_id
@@ -89,15 +97,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     director = C4Director(
         config[CONF_HOST], director_token_dict[CONF_TOKEN], director_session
     )
-    entry_data[CONF_DIRECTOR] = director
 
     controller_href = (await call_c4_api_retry(account.getAccountControllers))["href"]
-    entry_data[CONF_DIRECTOR_SW_VERSION] = await call_c4_api_retry(
+    director_sw_version = await call_c4_api_retry(
         account.getControllerOSVersion, controller_href
     )
 
     _, model, mac_address = controller_unique_id.split("_", 3)
-    entry_data[CONF_DIRECTOR_MODEL] = model.upper()
+    director_model = model.upper()
 
     device_registry = dr.async_get(hass)
     device_registry.async_get_or_create(
@@ -106,57 +113,60 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         connections={(dr.CONNECTION_NETWORK_MAC, mac_address)},
         manufacturer="Control4",
         name=controller_unique_id,
-        model=entry_data[CONF_DIRECTOR_MODEL],
-        sw_version=entry_data[CONF_DIRECTOR_SW_VERSION],
+        model=director_model,
+        sw_version=director_sw_version,
     )
 
     # Store all items found on controller for platforms to use
-    director_all_items = await director.getAllItemInfo()
-    director_all_items = json.loads(director_all_items)
-    entry_data[CONF_DIRECTOR_ALL_ITEMS] = director_all_items
+    director_all_items: list[dict[str, Any]] = json.loads(
+        await director.getAllItemInfo()
+    )
 
     # Check if OS version is 3 or higher to get UI configuration
-    entry_data[CONF_UI_CONFIGURATION] = None
-    if int(entry_data[CONF_DIRECTOR_SW_VERSION].split(".")[0]) >= 3:
-        entry_data[CONF_UI_CONFIGURATION] = json.loads(
-            await director.getUiConfiguration()
-        )
+    ui_configuration: dict[str, Any] | None = None
+    if int(director_sw_version.split(".")[0]) >= 3:
+        ui_configuration = json.loads(await director.getUiConfiguration())
 
     # Load options from config entry
-    entry_data[CONF_SCAN_INTERVAL] = entry.options.get(
-        CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL
+    scan_interval: int = entry.options.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
+
+    entry.runtime_data = Control4RuntimeData(
+        account=account,
+        controller_unique_id=controller_unique_id,
+        director=director,
+        director_all_items=director_all_items,
+        director_model=director_model,
+        director_sw_version=director_sw_version,
+        scan_interval=scan_interval,
+        ui_configuration=ui_configuration,
     )
 
-    entry_data[CONF_CONFIG_LISTENER] = entry.add_update_listener(update_listener)
+    entry.async_on_unload(entry.add_update_listener(update_listener))
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
-async def update_listener(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
+async def update_listener(
+    hass: HomeAssistant, config_entry: Control4ConfigEntry
+) -> None:
     """Update when config_entry options update."""
     _LOGGER.debug("Config entry was updated, rerunning setup")
     await hass.config_entries.async_reload(config_entry.entry_id)
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: Control4ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    hass.data[DOMAIN][entry.entry_id][CONF_CONFIG_LISTENER]()
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-        _LOGGER.debug("Unloaded entry for %s", entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
-async def get_items_of_category(hass: HomeAssistant, entry: ConfigEntry, category: str):
+async def get_items_of_category(
+    hass: HomeAssistant, entry: Control4ConfigEntry, category: str
+):
     """Return a list of all Control4 items with the specified category."""
-    director_all_items = hass.data[DOMAIN][entry.entry_id][CONF_DIRECTOR_ALL_ITEMS]
     return [
         item
-        for item in director_all_items
+        for item in entry.runtime_data.director_all_items
         if "categories" in item and category in item["categories"]
     ]

homeassistant/components/control4/config_flow.py

@@ -11,12 +11,7 @@
 from pyControl4.error_handling import NotFound, Unauthorized
 import voluptuous as vol
 
-from homeassistant.config_entries import (
-    ConfigEntry,
-    ConfigFlow,
-    ConfigFlowResult,
-    OptionsFlow,
-)
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult, OptionsFlow
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -28,6 +23,7 @@
 from homeassistant.helpers import aiohttp_client, config_validation as cv
 from homeassistant.helpers.device_registry import format_mac
 
+from . import Control4ConfigEntry
 from .const import (
     CONF_CONTROLLER_UNIQUE_ID,
     DEFAULT_SCAN_INTERVAL,
@@ -151,7 +147,7 @@ async def async_step_user(
     @staticmethod
     @callback
     def async_get_options_flow(
-        config_entry: ConfigEntry,
+        config_entry: Control4ConfigEntry,
     ) -> OptionsFlowHandler:
         """Get the options flow for this handler."""
         return OptionsFlowHandler()

homeassistant/components/control4/const.py

@@ -7,14 +7,6 @@
 
 API_RETRY_TIMES = 5
 
-CONF_ACCOUNT = "account"
-CONF_DIRECTOR = "director"
-CONF_DIRECTOR_SW_VERSION = "director_sw_version"
-CONF_DIRECTOR_MODEL = "director_model"
-CONF_DIRECTOR_ALL_ITEMS = "director_all_items"
-CONF_UI_CONFIGURATION = "ui_configuration"
 CONF_CONTROLLER_UNIQUE_ID = "controller_unique_id"
 
-CONF_CONFIG_LISTENER = "config_listener"
-
 CONTROL4_ENTITY_TYPE = 7

homeassistant/components/control4/director_utils.py

@@ -8,21 +8,21 @@
 from pyControl4.director import C4Director
 from pyControl4.error_handling import BadToken
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_TOKEN, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import aiohttp_client
 
-from .const import CONF_ACCOUNT, CONF_CONTROLLER_UNIQUE_ID, CONF_DIRECTOR, DOMAIN
+from . import Control4ConfigEntry
+from .const import CONF_CONTROLLER_UNIQUE_ID
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def _update_variables_for_config_entry(
-    hass: HomeAssistant, entry: ConfigEntry, variable_names: set[str]
+    hass: HomeAssistant, entry: Control4ConfigEntry, variable_names: set[str]
 ) -> dict[int, dict[str, Any]]:
     """Retrieve data from the Control4 director."""
-    director: C4Director = hass.data[DOMAIN][entry.entry_id][CONF_DIRECTOR]
+    director = entry.runtime_data.director
     data = await director.getAllItemVariableValue(variable_names)
     result_dict: defaultdict[int, dict[str, Any]] = defaultdict(dict)
     for item in data:
@@ -31,7 +31,7 @@ async def _update_variables_for_config_entry(
 
 
 async def update_variables_for_config_entry(
-    hass: HomeAssistant, entry: ConfigEntry, variable_names: set[str]
+    hass: HomeAssistant, entry: Control4ConfigEntry, variable_names: set[str]
 ) -> dict[int, dict[str, Any]]:
     """Try to Retrieve data from the Control4 director for update_coordinator."""
     try:
@@ -42,8 +42,8 @@ async def update_variables_for_config_entry(
         return await _update_variables_for_config_entry(hass, entry, variable_names)
 
 
-async def refresh_tokens(hass: HomeAssistant, entry: ConfigEntry):
-    """Store updated authentication and director tokens in hass.data."""
+async def refresh_tokens(hass: HomeAssistant, entry: Control4ConfigEntry):
+    """Store updated authentication and director tokens in runtime_data."""
     config = entry.data
     account_session = aiohttp_client.async_get_clientsession(hass)
 
@@ -59,6 +59,5 @@ async def refresh_tokens(hass: HomeAssistant, entry: ConfigEntry):
     )
 
     _LOGGER.debug("Saving new tokens in hass data")
-    entry_data = hass.data[DOMAIN][entry.entry_id]
-    entry_data[CONF_ACCOUNT] = account
-    entry_data[CONF_DIRECTOR] = director
+    entry.runtime_data.account = account
+    entry.runtime_data.director = director

homeassistant/components/control4/entity.py

@@ -10,15 +10,16 @@
     DataUpdateCoordinator,
 )
 
-from .const import CONF_CONTROLLER_UNIQUE_ID, DOMAIN
+from . import Control4RuntimeData
+from .const import DOMAIN
 
 
 class Control4Entity(CoordinatorEntity[Any]):
     """Base entity for Control4."""
 
     def __init__(
         self,
-        entry_data: dict,
+        runtime_data: Control4RuntimeData,
         coordinator: DataUpdateCoordinator[Any],
         name: str | None,
         idx: int,
@@ -29,11 +30,11 @@ def __init__(
     ) -> None:
         """Initialize a Control4 entity."""
         super().__init__(coordinator)
-        self.entry_data = entry_data
+        self.runtime_data = runtime_data
         self._attr_name = name
         self._attr_unique_id = str(idx)
         self._idx = idx
-        self._controller_unique_id = entry_data[CONF_CONTROLLER_UNIQUE_ID]
+        self._controller_unique_id = runtime_data.controller_unique_id
         self._device_name = device_name
         self._device_manufacturer = device_manufacturer
         self._device_model = device_model