Migrate elkm1 to use a dataclass for integration data

homeassistant/components/elkm1/__init__.py

@@ -2,11 +2,12 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Iterable
 from enum import Enum
 import logging
 import re
 from types import MappingProxyType
-from typing import Any, cast
+from typing import Any
 
 from elkm1_lib.elements import Element
 from elkm1_lib.elk import Elk
@@ -65,6 +66,7 @@ from .discovery import (
     async_trigger_discovery,
     async_update_entry_from_discovery,
 )
+from .models import ELKM1Data
 
 SYNC_TIMEOUT = 120
 
@@ -303,14 +305,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     else:
         temperature_unit = UnitOfTemperature.FAHRENHEIT
     config["temperature_unit"] = temperature_unit
-    hass.data[DOMAIN][entry.entry_id] = {
-        "elk": elk,
-        "prefix": conf[CONF_PREFIX],
-        "mac": entry.unique_id,
-        "auto_configure": conf[CONF_AUTO_CONFIGURE],
-        "config": config,
-        "keypads": {},
-    }
+    prefix: str = conf[CONF_PREFIX]
+    auto_configure: bool = conf[CONF_AUTO_CONFIGURE]
+    hass.data[DOMAIN][entry.entry_id] = ELKM1Data(
+        elk=elk,
+        prefix=prefix,
+        mac=entry.unique_id,
+        auto_configure=auto_configure,
+        config=config,
+        keypads={},
+    )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -326,21 +330,23 @@ def _included(ranges: list[tuple[int, int]], set_to: bool, values: list[bool]) -
 
 def _find_elk_by_prefix(hass: HomeAssistant, prefix: str) -> Elk | None:
     """Search all config entries for a given prefix."""
-    for entry_id in hass.data[DOMAIN]:
-        if hass.data[DOMAIN][entry_id]["prefix"] == prefix:
-            return cast(Elk, hass.data[DOMAIN][entry_id]["elk"])
+    all_elk: dict[str, ELKM1Data] = hass.data[DOMAIN]
+    for elk_data in all_elk.values():
+        if elk_data.prefix == prefix:
+            return elk_data.elk
     return None
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    all_elk: dict[str, ELKM1Data] = hass.data[DOMAIN]
 
     # disconnect cleanly
-    hass.data[DOMAIN][entry.entry_id]["elk"].disconnect()
+    all_elk[entry.entry_id].elk.disconnect()
 
     if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
+        all_elk.pop(entry.entry_id)
 
     return unload_ok
 
@@ -421,19 +427,19 @@ def _create_elk_services(hass: HomeAssistant) -> None:
 
 
 def create_elk_entities(
-    elk_data: dict[str, Any],
-    elk_elements: list[Element],
+    elk_data: ELKM1Data,
+    elk_elements: Iterable[Element],
     element_type: str,
     class_: Any,
     entities: list[ElkEntity],
 ) -> list[ElkEntity] | None:
     """Create the ElkM1 devices of a particular class."""
-    auto_configure = elk_data["auto_configure"]
+    auto_configure = elk_data.auto_configure
 
-    if not auto_configure and not elk_data["config"][element_type]["enabled"]:
+    if not auto_configure and not elk_data.config[element_type]["enabled"]:
         return None
 
-    elk = elk_data["elk"]
+    elk = elk_data.elk
     _LOGGER.debug("Creating elk entities for %s", elk)
 
     for element in elk_elements:
@@ -441,7 +447,7 @@ def create_elk_entities(
             if not element.configured:
                 continue
         # Only check the included list if auto configure is not
-        elif not elk_data["config"][element_type]["included"][element.index]:
+        elif not elk_data.config[element_type]["included"][element.index]:
             continue
 
         entities.append(class_(element, elk, elk_data))
@@ -454,13 +460,13 @@ class ElkEntity(Entity):
     _attr_has_entity_name = True
     _attr_should_poll = False
 
-    def __init__(self, element: Element, elk: Elk, elk_data: dict[str, Any]) -> None:
+    def __init__(self, element: Element, elk: Elk, elk_data: ELKM1Data) -> None:
         """Initialize the base of all Elk devices."""
         self._elk = elk
         self._element = element
-        self._mac = elk_data["mac"]
-        self._prefix = elk_data["prefix"]
-        self._temperature_unit: str = elk_data["config"]["temperature_unit"]
+        self._mac = elk_data.mac
+        self._prefix = elk_data.prefix
+        self._temperature_unit: str = elk_data.config["temperature_unit"]
         # unique_id starts with elkm1_ iff there is no prefix
         # it starts with elkm1m_{prefix} iff there is a prefix
         # this is to avoid a conflict between
@@ -496,9 +502,7 @@ class ElkEntity(Entity):
 
     def initial_attrs(self) -> dict[str, Any]:
         """Return the underlying element's attributes as a dict."""
-        attrs = {}
-        attrs["index"] = self._element.index + 1
-        return attrs
+        return {"index": self._element.index + 1}
 
     def _element_changed(self, element: Element, changeset: dict[str, Any]) -> None:
         pass

homeassistant/components/elkm1/alarm_control_panel.py

@@ -40,6 +40,7 @@ from .const import (
     DOMAIN,
     ELK_USER_CODE_SERVICE_SCHEMA,
 )
+from .models import ELKM1Data
 
 DISPLAY_MESSAGE_SERVICE_SCHEMA = {
     vol.Optional("clear", default=2): vol.All(vol.Coerce(int), vol.In([0, 1, 2])),
@@ -65,8 +66,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the ElkM1 alarm platform."""
-    elk_data = hass.data[DOMAIN][config_entry.entry_id]
-    elk = elk_data["elk"]
+
+    elk_data: ELKM1Data = hass.data[DOMAIN][config_entry.entry_id]
+    elk = elk_data.elk
     entities: list[ElkEntity] = []
     create_elk_entities(elk_data, elk.areas, "area", ElkArea, entities)
     async_add_entities(entities)
@@ -115,7 +117,7 @@ class ElkArea(ElkAttachedEntity, AlarmControlPanelEntity, RestoreEntity):
     )
     _element: Area
 
-    def __init__(self, element: Element, elk: Elk, elk_data: dict[str, Any]) -> None:
+    def __init__(self, element: Element, elk: Elk, elk_data: ELKM1Data) -> None:
         """Initialize Area as Alarm Control Panel."""
         super().__init__(element, elk, elk_data)
         self._elk = elk

homeassistant/components/elkm1/binary_sensor.py

@@ -14,6 +14,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import ElkAttachedEntity, ElkEntity
 from .const import DOMAIN
+from .models import ELKM1Data
 
 
 async def async_setup_entry(
@@ -22,21 +23,20 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Create the Elk-M1 sensor platform."""
-
-    elk_data = hass.data[DOMAIN][config_entry.entry_id]
-    auto_configure = elk_data["auto_configure"]
-    elk = elk_data["elk"]
+    elk_data: ELKM1Data = hass.data[DOMAIN][config_entry.entry_id]
+    elk = elk_data.elk
+    auto_configure = elk_data.auto_configure
 
     entities: list[ElkEntity] = []
     for element in elk.zones:
         # Don't create binary sensors for zones that are analog
-        if element.definition in {ZoneType.TEMPERATURE, ZoneType.ANALOG_ZONE}:
+        if element.definition in {ZoneType.TEMPERATURE, ZoneType.ANALOG_ZONE}:  # type: ignore[attr-defined]
             continue
 
         if auto_configure:
             if not element.configured:
                 continue
-        elif not elk_data["config"]["zone"]["included"][element.index]:
+        elif not elk_data.config["zone"]["included"][element.index]:
             continue
 
         entities.append(ElkBinarySensor(element, elk, elk_data))

homeassistant/components/elkm1/climate.py

@@ -23,6 +23,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import ElkEntity, create_elk_entities
 from .const import DOMAIN
+from .models import ELKM1Data
 
 SUPPORT_HVAC = [
     HVACMode.OFF,
@@ -61,9 +62,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Create the Elk-M1 thermostat platform."""
-    elk_data = hass.data[DOMAIN][config_entry.entry_id]
+    elk_data: ELKM1Data = hass.data[DOMAIN][config_entry.entry_id]
+    elk = elk_data.elk
     entities: list[ElkEntity] = []
-    elk = elk_data["elk"]
     create_elk_entities(
         elk_data, elk.thermostats, "thermostat", ElkThermostat, entities
     )

homeassistant/components/elkm1/light.py

@@ -14,6 +14,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import ElkEntity, create_elk_entities
 from .const import DOMAIN
+from .models import ELKM1Data
 
 
 async def async_setup_entry(
@@ -22,9 +23,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Elk light platform."""
-    elk_data = hass.data[DOMAIN][config_entry.entry_id]
+    elk_data: ELKM1Data = hass.data[DOMAIN][config_entry.entry_id]
+    elk = elk_data.elk
     entities: list[ElkEntity] = []
-    elk = elk_data["elk"]
     create_elk_entities(elk_data, elk.lights, "plc", ElkLight, entities)
     async_add_entities(entities)
 
@@ -36,7 +37,7 @@ class ElkLight(ElkEntity, LightEntity):
     _attr_supported_color_modes = {ColorMode.BRIGHTNESS}
     _element: Light
 
-    def __init__(self, element: Element, elk: Elk, elk_data: dict[str, Any]) -> None:
+    def __init__(self, element: Element, elk: Elk, elk_data: ELKM1Data) -> None:
         """Initialize the Elk light."""
         super().__init__(element, elk, elk_data)
         self._brightness = self._element.status

homeassistant/components/elkm1/models.py

@@ -0,0 +1,19 @@
+"""The elkm1 integration models."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from elkm1_lib import Elk
+
+
+@dataclass
+class ELKM1Data:
+    """Data for the elkm1 integration."""
+
+    elk: Elk
+    prefix: str
+    mac: str | None
+    auto_configure: bool
+    config: dict[str, Any]
+    keypads: dict[str, Any]

homeassistant/components/elkm1/switch.py

@@ -12,6 +12,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import ElkAttachedEntity, ElkEntity, create_elk_entities
 from .const import DOMAIN
+from .models import ELKM1Data
 
 
 async def async_setup_entry(
@@ -20,9 +21,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Create the Elk-M1 switch platform."""
-    elk_data = hass.data[DOMAIN][config_entry.entry_id]
+    elk_data: ELKM1Data = hass.data[DOMAIN][config_entry.entry_id]
+    elk = elk_data.elk
     entities: list[ElkEntity] = []
-    elk = elk_data["elk"]
     create_elk_entities(elk_data, elk.outputs, "output", ElkOutput, entities)
     async_add_entities(entities)