Store runtime data inside ConfigEntry

custom_components/vantage/__init__.py

@@ -11,11 +11,7 @@
 )
 from aiovantage.objects import Master
 
-from homeassistant.config_entries import (
-    ConfigEntry,
-    ConfigEntryAuthFailed,
-    ConfigEntryNotReady,
-)
+from homeassistant.config_entries import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -26,7 +22,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.util.ssl import get_default_no_verify_context
 
-from .const import DOMAIN
+from .config_entry import VantageConfigEntry, VantageData
 from .device import async_setup_devices
 from .entity import async_cleanup_entities
 from .events import async_setup_events
@@ -51,7 +47,7 @@
 Vantage.set_ssl_context_factory(get_default_no_verify_context)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: VantageConfigEntry) -> bool:
     """Set up Vantage integration from a config entry."""
     # Create a Vantage client
     vantage = Vantage(
@@ -61,8 +57,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         ssl=entry.data.get(CONF_SSL, True),
     )
 
-    # Store the client in the hass data store
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = vantage
+    # Store the client in the config entry's runtime data
+    entry.runtime_data = VantageData(client=vantage)
 
     try:
         # Initialize and fetch all objects
@@ -118,10 +114,10 @@ async def handle_system_program_event(
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: VantageConfigEntry) -> bool:
     """Unload a config entry."""
-    vantage: Vantage = hass.data[DOMAIN].pop(entry.entry_id, None)
-    if vantage:
-        vantage.close()
+    # Close the Vantage client connection
+    entry.runtime_data.client.close()
 
+    # Unload all platforms
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

custom_components/vantage/binary_sensor.py

@@ -2,23 +2,23 @@
 
 from functools import partial
 
-from aiovantage import Vantage
 from aiovantage.objects import DryContact
 
 from homeassistant.components.binary_sensor import BinarySensorEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from .config_entry import VantageConfigEntry
 from .entity import VantageEntity, async_register_vantage_objects
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: VantageConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Vantage binary sensor entities from a config entry."""
-    vantage: Vantage = hass.data[DOMAIN][entry.entry_id]
+    vantage = entry.runtime_data.client
     register_items = partial(
         async_register_vantage_objects, hass, entry, async_add_entities
     )

custom_components/vantage/climate.py

@@ -3,7 +3,7 @@
 import functools
 from typing import Any
 
-from aiovantage import Vantage, VantageEvent
+from aiovantage import VantageEvent
 from aiovantage.objects import Thermostat
 
 from homeassistant.components.climate import (
@@ -16,12 +16,12 @@
     HVACAction,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, LOGGER
+from .config_entry import VantageConfigEntry
+from .const import LOGGER
 from .entity import VantageEntity, async_register_vantage_objects
 
 # Set up the min/max temperature range for the thermostat
@@ -49,10 +49,12 @@
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: VantageConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Vantage cover entities from config entry."""
-    vantage: Vantage = hass.data[DOMAIN][entry.entry_id]
+    vantage = entry.runtime_data.client
     register_items = functools.partial(
         async_register_vantage_objects, hass, entry, async_add_entities
     )

custom_components/vantage/config_entry.py

@@ -0,0 +1,16 @@
+"""Vantage config entry."""
+
+from dataclasses import dataclass
+
+from homeassistant.config_entries import ConfigEntry
+
+from aiovantage import Vantage
+
+type VantageConfigEntry = ConfigEntry[VantageData]
+
+
+@dataclass
+class VantageData:
+    """Data for a Vantage config entry."""
+
+    client: Vantage

custom_components/vantage/config_flow.py

@@ -14,9 +14,10 @@
 
 from homeassistant import config_entries
 from homeassistant.components import zeroconf
-from homeassistant.config_entries import ConfigEntry, ConfigFlowResult
+from homeassistant.config_entries import ConfigFlowResult
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_SSL, CONF_USERNAME
 
+from .config_entry import VantageConfigEntry
 from .const import DOMAIN
 
 USER_SCHEMA = vol.Schema(
@@ -43,7 +44,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     controller: VantageControllerDetails | None = None
     username: str | None = None
     password: str | None = None
-    reauth_entry: ConfigEntry | None = None
+    reauth_entry: VantageConfigEntry | None = None
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None

custom_components/vantage/cover.py

@@ -3,7 +3,6 @@
 import functools
 from typing import Any
 
-from aiovantage import Vantage
 from aiovantage.objects import Blind, BlindGroup
 
 from homeassistant.components.cover import (
@@ -12,19 +11,20 @@
     CoverEntityFeature,
     ATTR_POSITION,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from .config_entry import VantageConfigEntry
 from .entity import VantageEntity, async_register_vantage_objects
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: VantageConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Vantage cover entities from config entry."""
-    vantage: Vantage = hass.data[DOMAIN][entry.entry_id]
+    vantage = entry.runtime_data.client
     register_items = functools.partial(
         async_register_vantage_objects, hass, entry, async_add_entities
     )

custom_components/vantage/device.py

@@ -12,19 +12,19 @@
     SystemObject,
 )
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo
 
+from .config_entry import VantageConfigEntry
 from .const import DOMAIN
 
 T = TypeVar("T", bound=SystemObject)
 
 
-def async_setup_devices(hass: HomeAssistant, entry: ConfigEntry) -> None:
+def async_setup_devices(hass: HomeAssistant, entry: VantageConfigEntry) -> None:
     """Set up Vantage devices in the device registry."""
-    vantage: Vantage = hass.data[DOMAIN][entry.entry_id]
+    vantage = entry.runtime_data.client
     dev_reg = dr.async_get(hass)
 
     @callback

custom_components/vantage/entity.py

@@ -13,13 +13,13 @@
 )
 from aiovantage.objects import GMem, SystemObject
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.entity import DeviceInfo, Entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from .config_entry import VantageConfigEntry
 from .const import DOMAIN
 from .device import vantage_device_info
 
@@ -32,14 +32,14 @@
 
 def async_register_vantage_objects(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: VantageConfigEntry,
     async_add_entities: AddEntitiesCallback,
     controller: BaseController[Any],
     entity_class: type["VantageEntity[Any]"],
     object_filter: Callable[[SystemObjectT], bool] | None = None,
 ) -> None:
     """Add entities to HA from a Vantage controller, add a callback for new entities."""
-    vantage: Vantage = hass.data[DOMAIN][config_entry.entry_id]
+    vantage = config_entry.runtime_data.client
 
     # Add all current objects in the controller that match the filter
     objects = controller.filter(object_filter) if object_filter else controller
@@ -57,9 +57,9 @@ def async_add_entity(_type: VantageEvent, obj: SystemObjectT, _data: Any) -> Non
     )
 
 
-def async_cleanup_entities(hass: HomeAssistant, entry: ConfigEntry) -> None:
+def async_cleanup_entities(hass: HomeAssistant, entry: VantageConfigEntry) -> None:
     """Remove entities from HA that are no longer in the Vantage controller."""
-    vantage: Vantage = hass.data[DOMAIN][entry.entry_id]
+    vantage = entry.runtime_data.client
     ent_reg = er.async_get(hass)
     for entity in er.async_entries_for_config_entry(ent_reg, entry.entry_id):
         # Entity IDs always start with the object ID, followed by an optional suffix
@@ -77,7 +77,7 @@ class VantageEntity(Generic[SystemObjectT], Entity):
     def __init__(
         self,
         client: Vantage,
-        config_entry: ConfigEntry,
+        config_entry: VantageConfigEntry,
         controller: BaseController[Any],
         obj: SystemObjectT,
     ):

custom_components/vantage/events.py

@@ -2,14 +2,13 @@
 
 from typing import Any
 
-from aiovantage import Vantage, VantageEvent
+from aiovantage import VantageEvent
 from aiovantage.objects import Button, Task
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
+from .config_entry import VantageConfigEntry
 from .const import (
-    DOMAIN,
     EVENT_BUTTON_PRESSED,
     EVENT_BUTTON_RELEASED,
     EVENT_TASK_STARTED,
@@ -18,9 +17,9 @@
 )
 
 
-def async_setup_events(hass: HomeAssistant, entry: ConfigEntry) -> None:
+def async_setup_events(hass: HomeAssistant, entry: VantageConfigEntry) -> None:
     """Set up Vantage events from a config entry."""
-    vantage: Vantage = hass.data[DOMAIN][entry.entry_id]
+    vantage = entry.runtime_data.client
 
     def handle_button_event(_event: VantageEvent, obj: Button, data: Any) -> None:
         """Handle button press/release events."""

custom_components/vantage/light.py

@@ -4,7 +4,6 @@
 import functools
 from typing import Any, TypeVar, cast
 
-from aiovantage import Vantage
 from aiovantage.objects import Load, LoadGroup, RGBLoadBase
 
 from homeassistant.components.light import (
@@ -18,14 +17,14 @@
     LightEntity,
     LightEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.color import brightness_to_value, value_to_brightness
 
-from .const import DOMAIN, LOGGER
+from .config_entry import VantageConfigEntry
+from .const import LOGGER
 from .entity import VantageEntity, async_register_vantage_objects
 
 # TypeVar for RGB/RGBW color tuples
@@ -36,10 +35,12 @@
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: VantageConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Vantage light entities from config entry."""
-    vantage: Vantage = hass.data[DOMAIN][entry.entry_id]
+    vantage = entry.runtime_data.client
     register_items = functools.partial(
         async_register_vantage_objects, hass, entry, async_add_entities
     )

custom_components/vantage/migrate.py

@@ -1,21 +1,21 @@
 """Migration functions for the Vantage integration."""
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers import entity_registry as er
 
+from .config_entry import VantageConfigEntry
 from .const import LOGGER
 
 
-async def async_migrate_data(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_migrate_data(hass: HomeAssistant, entry: VantageConfigEntry) -> None:
     """Run all Vantage data migrations."""
 
     async_delete_back_boxes(hass, entry)
     async_delete_serial_number_entities(hass, entry)
 
 
-def async_delete_back_boxes(hass: HomeAssistant, entry: ConfigEntry) -> None:
+def async_delete_back_boxes(hass: HomeAssistant, entry: VantageConfigEntry) -> None:
     """Delete back boxes from the device registry."""
     dev_reg = dr.async_get(hass)
 
@@ -33,7 +33,7 @@ def async_delete_back_boxes(hass: HomeAssistant, entry: ConfigEntry) -> None:
 
 
 def async_delete_serial_number_entities(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: VantageConfigEntry
 ) -> None:
     """Delete serial number entities from the entity registry."""
     ent_reg = er.async_get(hass)