Use EntityDescription - toon (#55035)

homeassistant/components/toon/switch.py

@@ -1,4 +1,7 @@
 """Support for Toon switches."""
+from __future__ import annotations
+
+from dataclasses import dataclass
 from typing import Any
 
 from toonapi import (
@@ -8,21 +11,14 @@ from toonapi import (
     PROGRAM_STATE_ON,
 )
 
-from homeassistant.components.switch import SwitchEntity
+from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import (
-    ATTR_ICON,
-    ATTR_MEASUREMENT,
-    ATTR_NAME,
-    ATTR_SECTION,
-    DOMAIN,
-    SWITCH_ENTITIES,
-)
+from .const import DOMAIN
 from .coordinator import ToonDataUpdateCoordinator
 from .helpers import toon_exception_handler
-from .models import ToonDisplayDeviceEntity, ToonEntity
+from .models import ToonDisplayDeviceEntity, ToonEntity, ToonRequiredKeysMixin
 
 
 async def async_setup_entry(
@@ -32,39 +28,38 @@ async def async_setup_entry(
     coordinator = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
-        [ToonProgramSwitch(coordinator), ToonHolidayModeSwitch(coordinator)]
+        [description.cls(coordinator, description) for description in SWITCH_ENTITIES]
     )
 
 
 class ToonSwitch(ToonEntity, SwitchEntity):
     """Defines an Toon switch."""
 
-    def __init__(self, coordinator: ToonDataUpdateCoordinator, *, key: str) -> None:
+    entity_description: ToonSwitchEntityDescription
+
+    def __init__(
+        self,
+        coordinator: ToonDataUpdateCoordinator,
+        description: ToonSwitchEntityDescription,
+    ) -> None:
         """Initialize the Toon switch."""
-        self.key = key
+        self.entity_description = description
         super().__init__(coordinator)
 
-        switch = SWITCH_ENTITIES[key]
-        self._attr_icon = switch[ATTR_ICON]
-        self._attr_name = switch[ATTR_NAME]
-        self._attr_unique_id = f"{coordinator.data.agreement.agreement_id}_{key}"
+        self._attr_unique_id = (
+            f"{coordinator.data.agreement.agreement_id}_{description.key}"
+        )
 
     @property
     def is_on(self) -> bool:
         """Return the status of the binary sensor."""
-        section = getattr(
-            self.coordinator.data, SWITCH_ENTITIES[self.key][ATTR_SECTION]
-        )
-        return getattr(section, SWITCH_ENTITIES[self.key][ATTR_MEASUREMENT])
+        section = getattr(self.coordinator.data, self.entity_description.section)
+        return getattr(section, self.entity_description.measurement)
 
 
 class ToonProgramSwitch(ToonSwitch, ToonDisplayDeviceEntity):
     """Defines a Toon program switch."""
 
-    def __init__(self, coordinator: ToonDataUpdateCoordinator) -> None:
-        """Initialize the Toon program switch."""
-        super().__init__(coordinator, key="thermostat_program")
-
     @toon_exception_handler
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the Toon program switch."""
@@ -83,10 +78,6 @@ class ToonProgramSwitch(ToonSwitch, ToonDisplayDeviceEntity):
 class ToonHolidayModeSwitch(ToonSwitch, ToonDisplayDeviceEntity):
     """Defines a Toon Holiday mode switch."""
 
-    def __init__(self, coordinator: ToonDataUpdateCoordinator) -> None:
-        """Initialize the Toon holiday switch."""
-        super().__init__(coordinator, key="thermostat_holiday_mode")
-
     @toon_exception_handler
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the Toon holiday mode switch."""
@@ -100,3 +91,35 @@ class ToonHolidayModeSwitch(ToonSwitch, ToonDisplayDeviceEntity):
         await self.coordinator.toon.set_active_state(
             ACTIVE_STATE_HOLIDAY, PROGRAM_STATE_OFF
         )
+
+
+@dataclass
+class ToonSwitchRequiredKeysMixin(ToonRequiredKeysMixin):
+    """Mixin for switch required keys."""
+
+    cls: type[ToonSwitch]
+
+
+@dataclass
+class ToonSwitchEntityDescription(SwitchEntityDescription, ToonSwitchRequiredKeysMixin):
+    """Describes Toon switch entity."""
+
+
+SWITCH_ENTITIES: tuple[ToonSwitchEntityDescription, ...] = (
+    ToonSwitchEntityDescription(
+        key="thermostat_holiday_mode",
+        name="Holiday Mode",
+        section="thermostat",
+        measurement="holiday_mode",
+        icon="mdi:airport",
+        cls=ToonHolidayModeSwitch,
+    ),
+    ToonSwitchEntityDescription(
+        key="thermostat_program",
+        name="Thermostat Program",
+        section="thermostat",
+        measurement="program",
+        icon="mdi:calendar-clock",
+        cls=ToonProgramSwitch,
+    ),
+)