Move Lutron entry data to typed class (#107256)

* Move Lutron entry data to typed class

* Move Lutron entry data to typed class

* Exclude file

* Update homeassistant/components/lutron/__init__.py

Co-authored-by: Jan-Philipp Benecke <github@bnck.me>

---------

Co-authored-by: Jan-Philipp Benecke <github@bnck.me>

.coveragerc

@@ -699,6 +699,7 @@ omit =
     homeassistant/components/lutron/__init__.py
     homeassistant/components/lutron/binary_sensor.py
     homeassistant/components/lutron/cover.py
+    homeassistant/components/lutron/entity.py
     homeassistant/components/lutron/light.py
     homeassistant/components/lutron/switch.py
     homeassistant/components/lutron_caseta/__init__.py

homeassistant/components/lutron/__init__.py

@@ -1,7 +1,8 @@
 """Component for interacting with a Lutron RadioRA 2 system."""
+from dataclasses import dataclass
 import logging
 
-from pylutron import Button, Lutron
+from pylutron import Button, Led, Lutron, OccupancyGroup, Output
 import voluptuous as vol
 
 from homeassistant import config_entries
@@ -32,10 +33,6 @@ PLATFORMS = [
 
 _LOGGER = logging.getLogger(__name__)
 
-LUTRON_BUTTONS = "lutron_buttons"
-LUTRON_CONTROLLER = "lutron_controller"
-LUTRON_DEVICES = "lutron_devices"
-
 # Attribute on events that indicates what action was taken with the button.
 ATTR_ACTION = "action"
 ATTR_FULL_ID = "full_id"
@@ -105,78 +102,6 @@ async def async_setup(hass: HomeAssistant, base_config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
-    """Set up the Lutron integration."""
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[LUTRON_BUTTONS] = []
-    hass.data[LUTRON_CONTROLLER] = None
-    hass.data[LUTRON_DEVICES] = {
-        "light": [],
-        "cover": [],
-        "switch": [],
-        "scene": [],
-        "binary_sensor": [],
-    }
-    host = config_entry.data[CONF_HOST]
-    uid = config_entry.data[CONF_USERNAME]
-    pwd = config_entry.data[CONF_PASSWORD]
-
-    def _load_db() -> bool:
-        hass.data[LUTRON_CONTROLLER].load_xml_db()
-        return True
-
-    hass.data[LUTRON_CONTROLLER] = Lutron(host, uid, pwd)
-    await hass.async_add_executor_job(_load_db)
-    hass.data[LUTRON_CONTROLLER].connect()
-    _LOGGER.info("Connected to main repeater at %s", host)
-
-    # Sort our devices into types
-    _LOGGER.debug("Start adding devices")
-    for area in hass.data[LUTRON_CONTROLLER].areas:
-        _LOGGER.debug("Working on area %s", area.name)
-        for output in area.outputs:
-            _LOGGER.debug("Working on output %s", output.type)
-            if output.type == "SYSTEM_SHADE":
-                hass.data[LUTRON_DEVICES]["cover"].append((area.name, output))
-            elif output.is_dimmable:
-                hass.data[LUTRON_DEVICES]["light"].append((area.name, output))
-            else:
-                hass.data[LUTRON_DEVICES]["switch"].append((area.name, output))
-        for keypad in area.keypads:
-            for button in keypad.buttons:
-                # If the button has a function assigned to it, add it as a scene
-                if button.name != "Unknown Button" and button.button_type in (
-                    "SingleAction",
-                    "Toggle",
-                    "SingleSceneRaiseLower",
-                    "MasterRaiseLower",
-                ):
-                    # Associate an LED with a button if there is one
-                    led = next(
-                        (led for led in keypad.leds if led.number == button.number),
-                        None,
-                    )
-                    hass.data[LUTRON_DEVICES]["scene"].append(
-                        (area.name, keypad.name, button, led)
-                    )
-
-                hass.data[LUTRON_BUTTONS].append(
-                    LutronButton(hass, area.name, keypad, button)
-                )
-        if area.occupancy_group is not None:
-            hass.data[LUTRON_DEVICES]["binary_sensor"].append(
-                (area.name, area.occupancy_group)
-            )
-    await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
-
-    return True
-
-
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    """Clean up resources and entities associated with the integration."""
-    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-
 class LutronButton:
     """Representation of a button on a Lutron keypad.
 
@@ -227,3 +152,80 @@ class LutronButton:
                 ATTR_UUID: self._uuid,
             }
             self._hass.bus.fire(self._event, data)
+
+
+@dataclass(slots=True)
+class LutronData:
+    """Storage class for platform global data."""
+
+    client: Lutron
+    covers: list[tuple[str, Output]]
+    lights: list[tuple[str, Output]]
+    switches: list[tuple[str, Output]]
+    scenes: list[tuple[str, str, Button, Led]]
+    binary_sensors: list[tuple[str, OccupancyGroup]]
+    buttons: list[LutronButton]
+
+
+async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+    """Set up the Lutron integration."""
+
+    host = config_entry.data[CONF_HOST]
+    uid = config_entry.data[CONF_USERNAME]
+    pwd = config_entry.data[CONF_PASSWORD]
+
+    lutron_client = Lutron(host, uid, pwd)
+    await hass.async_add_executor_job(lutron_client.load_xml_db)
+    lutron_client.connect()
+    _LOGGER.info("Connected to main repeater at %s", host)
+
+    entry_data = LutronData(
+        client=lutron_client,
+        covers=[],
+        lights=[],
+        switches=[],
+        scenes=[],
+        binary_sensors=[],
+        buttons=[],
+    )
+    # Sort our devices into types
+    _LOGGER.debug("Start adding devices")
+    for area in lutron_client.areas:
+        _LOGGER.debug("Working on area %s", area.name)
+        for output in area.outputs:
+            _LOGGER.debug("Working on output %s", output.type)
+            if output.type == "SYSTEM_SHADE":
+                entry_data.covers.append((area.name, output))
+            elif output.is_dimmable:
+                entry_data.lights.append((area.name, output))
+            else:
+                entry_data.switches.append((area.name, output))
+        for keypad in area.keypads:
+            for button in keypad.buttons:
+                # If the button has a function assigned to it, add it as a scene
+                if button.name != "Unknown Button" and button.button_type in (
+                    "SingleAction",
+                    "Toggle",
+                    "SingleSceneRaiseLower",
+                    "MasterRaiseLower",
+                ):
+                    # Associate an LED with a button if there is one
+                    led = next(
+                        (led for led in keypad.leds if led.number == button.number),
+                        None,
+                    )
+                    entry_data.scenes.append((area.name, keypad.name, button, led))
+
+                entry_data.buttons.append(LutronButton(hass, area.name, keypad, button))
+        if area.occupancy_group is not None:
+            entry_data.binary_sensors.append((area.name, area.occupancy_group))
+    hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = entry_data
+
+    await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
+
+    return True
+
+
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Clean up resources and entities associated with the integration."""
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/lutron/binary_sensor.py

@@ -14,9 +14,8 @@ from homeassistant.components.binary_sensor import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.typing import DiscoveryInfoType
 
-from . import LUTRON_CONTROLLER, LUTRON_DEVICES
+from . import DOMAIN, LutronData
 from .entity import LutronDevice
 
 _LOGGER = logging.getLogger(__name__)
@@ -26,18 +25,20 @@ async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
     async_add_entities: AddEntitiesCallback,
-    discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
     """Set up the Lutron binary_sensor platform.
 
     Adds occupancy groups from the Main Repeater associated with the
     config_entry as binary_sensor entities.
     """
-    entities = []
+    entry_data: LutronData = hass.data[DOMAIN][config_entry.entry_id]
-    for area_name, device in hass.data[LUTRON_DEVICES]["binary_sensor"]:
+    async_add_entities(
-        entity = LutronOccupancySensor(area_name, device, hass.data[LUTRON_CONTROLLER])
+        [
-        entities.append(entity)
+            LutronOccupancySensor(area_name, device, entry_data.client)
-    async_add_entities(entities, True)
+            for area_name, device in entry_data.binary_sensors
+        ],
+        True,
+    )
 
 
 class LutronOccupancySensor(LutronDevice, BinarySensorEntity):

homeassistant/components/lutron/cover.py

@@ -14,7 +14,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LUTRON_CONTROLLER, LUTRON_DEVICES
+from . import DOMAIN, LutronData
 from .entity import LutronDevice
 
 _LOGGER = logging.getLogger(__name__)
@@ -30,11 +30,14 @@ async def async_setup_entry(
     Adds shades from the Main Repeater associated with the config_entry as
     cover entities.
     """
-    entities = []
+    entry_data: LutronData = hass.data[DOMAIN][config_entry.entry_id]
-    for area_name, device in hass.data[LUTRON_DEVICES]["cover"]:
+    async_add_entities(
-        entity = LutronCover(area_name, device, hass.data[LUTRON_CONTROLLER])
+        [
-        entities.append(entity)
+            LutronCover(area_name, device, entry_data.covers)
-    async_add_entities(entities, True)
+            for area_name, device in entry_data.covers
+        ],
+        True,
+    )
 
 
 class LutronCover(LutronDevice, CoverEntity):

homeassistant/components/lutron/light.py

@@ -9,7 +9,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LUTRON_CONTROLLER, LUTRON_DEVICES
+from . import DOMAIN, LutronData
 from .entity import LutronDevice
 
 
@@ -23,11 +23,14 @@ async def async_setup_entry(
     Adds dimmers from the Main Repeater associated with the config_entry as
     light entities.
     """
-    entities = []
+    entry_data: LutronData = hass.data[DOMAIN][config_entry.entry_id]
-    for area_name, device in hass.data[LUTRON_DEVICES]["light"]:
+    async_add_entities(
-        entity = LutronLight(area_name, device, hass.data[LUTRON_CONTROLLER])
+        [
-        entities.append(entity)
+            LutronLight(area_name, device, entry_data.client)
-    async_add_entities(entities, True)
+            for area_name, device in entry_data.lights
+        ],
+        True,
+    )
 
 
 def to_lutron_level(level):

homeassistant/components/lutron/scene.py

@@ -8,7 +8,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LUTRON_CONTROLLER, LUTRON_DEVICES
+from . import DOMAIN, LutronData
 from .entity import LutronDevice
 
 
@@ -22,14 +22,15 @@ async def async_setup_entry(
     Adds scenes from the Main Repeater associated with the config_entry as
     scene entities.
     """
-    entities = []
+    entry_data: LutronData = hass.data[DOMAIN][config_entry.entry_id]
-    for scene_data in hass.data[LUTRON_DEVICES]["scene"]:
+
-        (area_name, keypad_name, device, led) = scene_data
+    async_add_entities(
-        entity = LutronScene(
+        [
-            area_name, keypad_name, device, led, hass.data[LUTRON_CONTROLLER]
+            LutronScene(area_name, keypad_name, device, led, entry_data.client)
-        )
+            for area_name, keypad_name, device, led in entry_data.scenes
-        entities.append(entity)
+        ],
-    async_add_entities(entities, True)
+        True,
+    )
 
 
 class LutronScene(LutronDevice, Scene):

homeassistant/components/lutron/switch.py

@@ -9,7 +9,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LUTRON_CONTROLLER, LUTRON_DEVICES
+from . import DOMAIN, LutronData
 from .entity import LutronDevice
 
 
@@ -23,21 +23,19 @@ async def async_setup_entry(
     Adds switches from the Main Repeater associated with the config_entry as
     switch entities.
     """
-    entities = []
+    entry_data: LutronData = hass.data[DOMAIN][config_entry.entry_id]
+    entities: list[SwitchEntity] = []
 
     # Add Lutron Switches
-    for area_name, device in hass.data[LUTRON_DEVICES]["switch"]:
+    for area_name, device in entry_data.switches:
-        entity = LutronSwitch(area_name, device, hass.data[LUTRON_CONTROLLER])
+        entities.append(LutronSwitch(area_name, device, entry_data.client))
-        entities.append(entity)
 
     # Add the indicator LEDs for scenes (keypad buttons)
-    for scene_data in hass.data[LUTRON_DEVICES]["scene"]:
+    for area_name, keypad_name, scene, led in entry_data.scenes:
-        (area_name, keypad_name, scene, led) = scene_data
         if led is not None:
-            led = LutronLed(
+            entities.append(
-                area_name, keypad_name, scene, led, hass.data[LUTRON_CONTROLLER]
+                LutronLed(area_name, keypad_name, scene, led, entry_data.client)
             )
-            entities.append(led)
     async_add_entities(entities, True)