Use runtime_data in mystrom (home-assistant#148020)

Author: epenet

homeassistant/components/mystrom/__init__.py

@@ -9,13 +9,11 @@
 from pymystrom.exceptions import MyStromConnectionError
 from pymystrom.switch import MyStromSwitch
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
-from .const import DOMAIN
-from .models import MyStromData
+from .models import MyStromConfigEntry, MyStromData
 
 PLATFORMS_PLUGS = [Platform.SENSOR, Platform.SWITCH]
 PLATFORMS_BULB = [Platform.LIGHT]
@@ -41,7 +39,7 @@ def _get_mystrom_switch(host: str) -> MyStromSwitch:
     return MyStromSwitch(host)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: MyStromConfigEntry) -> bool:
     """Set up myStrom from a config entry."""
     host = entry.data[CONF_HOST]
     try:
@@ -73,7 +71,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         _LOGGER.error("Unsupported myStrom device type: %s", device_type)
         return False
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = MyStromData(
+    entry.runtime_data = MyStromData(
         device=device,
         info=info,
     )
@@ -82,15 +80,12 @@ 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: MyStromConfigEntry) -> bool:
     """Unload a config entry."""
-    device_type = hass.data[DOMAIN][entry.entry_id].info["type"]
+    device_type = entry.runtime_data.info["type"]
     platforms = []
     if device_type in [101, 106, 107, 120]:
         platforms.extend(PLATFORMS_PLUGS)
     elif device_type in [102, 105]:
         platforms.extend(PLATFORMS_BULB)
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, platforms):
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, platforms)

homeassistant/components/mystrom/light.py

@@ -15,12 +15,12 @@
     LightEntity,
     LightEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from .const import DOMAIN, MANUFACTURER
+from .models import MyStromConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,12 +32,12 @@
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: MyStromConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the myStrom entities."""
-    info = hass.data[DOMAIN][entry.entry_id].info
-    device = hass.data[DOMAIN][entry.entry_id].device
+    info = entry.runtime_data.info
+    device = entry.runtime_data.device
     async_add_entities([MyStromLight(device, entry.title, info["mac"])])
 
 

homeassistant/components/mystrom/models.py

@@ -6,6 +6,10 @@
 from pymystrom.bulb import MyStromBulb
 from pymystrom.switch import MyStromSwitch
 
+from homeassistant.config_entries import ConfigEntry
+
+type MyStromConfigEntry = ConfigEntry[MyStromData]
+
 
 @dataclass
 class MyStromData:

homeassistant/components/mystrom/sensor.py

@@ -13,13 +13,13 @@
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfPower, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from .const import DOMAIN, MANUFACTURER
+from .models import MyStromConfigEntry
 
 
 @dataclass(frozen=True)
@@ -56,11 +56,11 @@ class MyStromSwitchSensorEntityDescription(SensorEntityDescription):
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: MyStromConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the myStrom entities."""
-    device: MyStromSwitch = hass.data[DOMAIN][entry.entry_id].device
+    device: MyStromSwitch = entry.runtime_data.device
 
     async_add_entities(
         MyStromSwitchSensor(device, entry.title, description)

homeassistant/components/mystrom/switch.py

@@ -8,12 +8,12 @@
 from pymystrom.exceptions import MyStromConnectionError
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo, format_mac
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from .const import DOMAIN, MANUFACTURER
+from .models import MyStromConfigEntry
 
 DEFAULT_NAME = "myStrom Switch"
 
@@ -22,11 +22,11 @@
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: MyStromConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the myStrom entities."""
-    device = hass.data[DOMAIN][entry.entry_id].device
+    device = entry.runtime_data.device
     async_add_entities([MyStromSwitch(device, entry.title)])