Use runtime_data in hlk_sw16 (#144370)

homeassistant/components/hlk_sw16/__init__.py

@@ -3,6 +3,7 @@
 import logging
 
 from hlk_sw16 import create_hlk_sw16_connection
+from hlk_sw16.protocol import SW16Client
 import voluptuous as vol
 
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
@@ -24,9 +25,6 @@ _LOGGER = logging.getLogger(__name__)
 
 PLATFORMS = [Platform.SWITCH]
 
-DATA_DEVICE_REGISTER = "hlk_sw16_device_register"
-DATA_DEVICE_LISTENER = "hlk_sw16_device_listener"
-
 SWITCH_SCHEMA = vol.Schema({vol.Optional(CONF_NAME): cv.string})
 
 RELAY_ID = vol.All(
@@ -52,6 +50,8 @@ CONFIG_SCHEMA = vol.Schema(
     extra=vol.ALLOW_EXTRA,
 )
 
+type HlkConfigEntry = ConfigEntry[SW16Client]
+
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Component setup, do nothing."""
@@ -70,15 +70,12 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: HlkConfigEntry) -> bool:
     """Set up the HLK-SW16 switch."""
-    hass.data.setdefault(DOMAIN, {})
     host = entry.data[CONF_HOST]
     port = entry.data[CONF_PORT]
     address = f"{host}:{port}"
 
-    hass.data[DOMAIN][entry.entry_id] = {}
-
     @callback
     def disconnected():
         """Schedule reconnect after connection has been lost."""
@@ -106,7 +103,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         keep_alive_interval=DEFAULT_KEEP_ALIVE_INTERVAL,
     )
 
-    hass.data[DOMAIN][entry.entry_id][DATA_DEVICE_REGISTER] = client
+    entry.runtime_data = client
 
     # Load entities
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
@@ -116,14 +113,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: HlkConfigEntry) -> bool:
     """Unload a config entry."""
-    client = hass.data[DOMAIN][entry.entry_id].pop(DATA_DEVICE_REGISTER)
-    client.stop()
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        if hass.data[DOMAIN][entry.entry_id]:
-            hass.data[DOMAIN].pop(entry.entry_id)
-        if not hass.data[DOMAIN]:
-            hass.data.pop(DOMAIN)
-    return unload_ok
+    entry.runtime_data.stop()
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/hlk_sw16/entity.py

@@ -2,6 +2,8 @@
 
 import logging
 
+from hlk_sw16.protocol import SW16Client
+
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import Entity
@@ -17,12 +19,12 @@ class SW16Entity(Entity):
 
     _attr_should_poll = False
 
-    def __init__(self, device_port, entry_id, client):
+    def __init__(self, device_port: str, entry_id: str, client: SW16Client) -> None:
         """Initialize the device."""
         # HLK-SW16 specific attributes for every component type
         self._entry_id = entry_id
         self._device_port = device_port
-        self._is_on = None
+        self._is_on: bool | None = None
         self._client = client
         self._attr_name = device_port
         self._attr_unique_id = f"{self._entry_id}_{self._device_port}"

homeassistant/components/hlk_sw16/switch.py

@@ -1,22 +1,22 @@
 """Support for HLK-SW16 switches."""
 
+from __future__ import annotations
+
 from typing import Any
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import DATA_DEVICE_REGISTER
-from .const import DOMAIN
+from . import HlkConfigEntry
 from .entity import SW16Entity
 
 PARALLEL_UPDATES = 0
 
 
-def devices_from_entities(hass, entry):
+def devices_from_entities(entry: HlkConfigEntry) -> list[SW16Switch]:
     """Parse configuration and add HLK-SW16 switch devices."""
-    device_client = hass.data[DOMAIN][entry.entry_id][DATA_DEVICE_REGISTER]
+    device_client = entry.runtime_data
     devices = []
     for i in range(16):
         device_port = f"{i:01x}"
@@ -27,18 +27,18 @@ def devices_from_entities(hass, entry):
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: HlkConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the HLK-SW16 platform."""
-    async_add_entities(devices_from_entities(hass, entry))
+    async_add_entities(devices_from_entities(entry))
 
 
 class SW16Switch(SW16Entity, SwitchEntity):
     """Representation of a HLK-SW16 switch."""
 
     @property
-    def is_on(self):
+    def is_on(self) -> bool | None:
         """Return true if device is on."""
         return self._is_on