diff --git a/homeassistant/components/mystrom/__init__.py b/homeassistant/components/mystrom/__init__.py index 09cd7b42da0..9094fc11e1c 100644 --- a/homeassistant/components/mystrom/__init__.py +++ b/homeassistant/components/mystrom/__init__.py @@ -9,13 +9,11 @@ from pymystrom.bulb import MyStromBulb 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) diff --git a/homeassistant/components/mystrom/light.py b/homeassistant/components/mystrom/light.py index 3942f601a20..67964d7d5b4 100644 --- a/homeassistant/components/mystrom/light.py +++ b/homeassistant/components/mystrom/light.py @@ -15,12 +15,12 @@ from homeassistant.components.light import ( 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 @@ EFFECT_SUNRISE = "sunrise" 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"])]) diff --git a/homeassistant/components/mystrom/models.py b/homeassistant/components/mystrom/models.py index 694a2f43df6..a96837070fd 100644 --- a/homeassistant/components/mystrom/models.py +++ b/homeassistant/components/mystrom/models.py @@ -6,6 +6,10 @@ from typing import Any from pymystrom.bulb import MyStromBulb from pymystrom.switch import MyStromSwitch +from homeassistant.config_entries import ConfigEntry + +type MyStromConfigEntry = ConfigEntry[MyStromData] + @dataclass class MyStromData: diff --git a/homeassistant/components/mystrom/sensor.py b/homeassistant/components/mystrom/sensor.py index bd5c9b923a2..251765d1658 100644 --- a/homeassistant/components/mystrom/sensor.py +++ b/homeassistant/components/mystrom/sensor.py @@ -13,13 +13,13 @@ from homeassistant.components.sensor import ( 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 @@ SENSOR_TYPES: tuple[MyStromSwitchSensorEntityDescription, ...] = ( 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) diff --git a/homeassistant/components/mystrom/switch.py b/homeassistant/components/mystrom/switch.py index f626656a4e3..860d2dff727 100644 --- a/homeassistant/components/mystrom/switch.py +++ b/homeassistant/components/mystrom/switch.py @@ -8,12 +8,12 @@ from typing import Any 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 @@ _LOGGER = logging.getLogger(__name__) 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)])