Refactor switch for vesync

homeassistant/components/vesync/common.py

@@ -12,6 +12,28 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+def rgetattr(obj, attr):
+    """Return a string in the form word.1.2.3 and return the item as 3. Note that this last value could be in a dict as well."""
+    _this_func = rgetattr
+    sp = attr.split(".", 1)
+    if len(sp) == 1:
+        left, right = sp[0], ""
+    else:
+        left, right = sp
+
+    if isinstance(obj, dict):
+        obj = obj.get(left)
+    elif hasattr(obj, left):
+        obj = getattr(obj, left)
+    else:
+        return None
+
+    if right:
+        obj = _this_func(obj, right)
+
+    return obj
+
+
 async def async_process_devices(
     hass: HomeAssistant, manager: VeSync
 ) -> dict[str, list[VeSyncBaseDevice]]:

homeassistant/components/vesync/strings.json

@@ -43,6 +43,11 @@
         "name": "Current voltage"
       }
     },
+    "switch": {
+      "on": {
+        "name": "On"
+      }
+    },
     "fan": {
       "vesync": {
         "state_attributes": {

homeassistant/components/vesync/switch.py

@@ -1,29 +1,54 @@
 """Support for VeSync switches."""
 
+from collections.abc import Callable
+from dataclasses import dataclass
 import logging
-from typing import Any
+from typing import Any, Final
 
 from pyvesync.vesyncbasedevice import VeSyncBaseDevice
-
-from homeassistant.components.switch import SwitchEntity
+from pyvesync.vesyncoutlet import VeSyncOutlet
+from pyvesync.vesyncswitch import VeSyncSwitch
+
+from homeassistant.components.switch import (
+    SwitchDeviceClass,
+    SwitchEntity,
+    SwitchEntityDescription,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DEV_TYPE_TO_HA, DOMAIN, VS_COORDINATOR, VS_DISCOVERY, VS_SWITCHES
+from .common import rgetattr
+from .const import DOMAIN, VS_COORDINATOR, VS_DISCOVERY, VS_SWITCHES
 from .coordinator import VeSyncDataCoordinator
 from .entity import VeSyncBaseEntity
 
 _LOGGER = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True, kw_only=True)
+class VeSyncSwitchEntityDescription(SwitchEntityDescription):
+    """A class that describes custom switch entities."""
+
+    is_on: Callable[[VeSyncBaseDevice], bool]
+
+
+SENSOR_DESCRIPTIONS: Final[tuple[VeSyncSwitchEntityDescription, ...]] = (
+    VeSyncSwitchEntityDescription(
+        key="device_status",
+        translation_key="on",
+        is_on=lambda device: device.device_status == "on",
+    ),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Set up switches."""
+    """Set up switch platform."""
 
     coordinator = hass.data[DOMAIN][VS_COORDINATOR]
 
@@ -46,57 +71,46 @@
     coordinator: VeSyncDataCoordinator,
 ):
     """Check if device is online and add entity."""
-    entities: list[VeSyncBaseSwitch] = []
-    for dev in devices:
-        if DEV_TYPE_TO_HA.get(dev.device_type) == "outlet":
-            entities.append(VeSyncSwitchHA(dev, coordinator))
-        elif DEV_TYPE_TO_HA.get(dev.device_type) == "switch":
-            entities.append(VeSyncLightSwitch(dev, coordinator))
-        else:
-            _LOGGER.warning(
-                "%s - Unknown device type - %s", dev.device_name, dev.device_type
-            )
-            continue
-
-    async_add_entities(entities, update_before_add=True)
-
+    async_add_entities(
+        (
+            VeSyncSwitchEntity(dev, description, coordinator)
+            for dev in devices
+            for description in SENSOR_DESCRIPTIONS
+            if rgetattr(dev, description.key) is not None
+        ),
+        update_before_add=True,
+    )
 
-class VeSyncBaseSwitch(VeSyncBaseEntity, SwitchEntity):
-    """Base class for VeSync switch Device Representations."""
 
-    _attr_name = None
+class VeSyncSwitchEntity(SwitchEntity, VeSyncBaseEntity):
+    """VeSync sensor class."""
 
-    def turn_on(self, **kwargs: Any) -> None:
-        """Turn the device on."""
-        self.device.turn_on()
+    def __init__(
+        self,
+        device: VeSyncBaseDevice,
+        description: VeSyncSwitchEntityDescription,
+        coordinator: VeSyncDataCoordinator,
+    ) -> None:
+        """Initialize the sensor."""
+        super().__init__(device, coordinator)
+        self.entity_description: VeSyncSwitchEntityDescription = description
+        if isinstance(self.device, VeSyncOutlet):
+            self._attr_device_class = SwitchDeviceClass.OUTLET
+        if isinstance(self.device, VeSyncSwitch):
+            self._attr_device_class = SwitchDeviceClass.SWITCH
+        self._attr_name = None
 
     @property
-    def is_on(self) -> bool:
-        """Return True if device is on."""
-        return self.device.device_status == "on"
+    def is_on(self) -> bool | None:
+        """Return the entity value to represent the entity state."""
+        if self.entity_description.is_on is not None:
+            return self.entity_description.is_on(self.device)
+        return None
 
     def turn_off(self, **kwargs: Any) -> None:
-        """Turn the device off."""
+        """Turn the entity off."""
         self.device.turn_off()
 
-
-class VeSyncSwitchHA(VeSyncBaseSwitch, SwitchEntity):
-    """Representation of a VeSync switch."""
-
-    def __init__(
-        self, plug: VeSyncBaseDevice, coordinator: VeSyncDataCoordinator
-    ) -> None:
-        """Initialize the VeSync switch device."""
-        super().__init__(plug, coordinator)
-        self.smartplug = plug
-
-
-class VeSyncLightSwitch(VeSyncBaseSwitch, SwitchEntity):
-    """Handle representation of VeSync Light Switch."""
-
-    def __init__(
-        self, switch: VeSyncBaseDevice, coordinator: VeSyncDataCoordinator
-    ) -> None:
-        """Initialize Light Switch device class."""
-        super().__init__(switch, coordinator)
-        self.switch = switch
+    def turn_on(self, **kwargs: Any) -> None:
+        """Turn the entity on."""
+        self.device.turn_on()

tests/components/vesync/snapshots/test_switch.ambr

@@ -300,13 +300,13 @@
       'name': None,
       'options': dict({
       }),
-      'original_device_class': None,
+      'original_device_class': <SwitchDeviceClass.OUTLET: 'outlet'>,
       'original_icon': None,
       'original_name': None,
       'platform': 'vesync',
       'previous_unique_id': None,
       'supported_features': 0,
-      'translation_key': None,
+      'translation_key': 'on',
       'unique_id': 'outlet',
       'unit_of_measurement': None,
     }),
@@ -315,6 +315,7 @@
 # name: test_switch_state[Outlet][switch.outlet]
   StateSnapshot({
     'attributes': ReadOnlyDict({
+      'device_class': 'outlet',
       'friendly_name': 'Outlet',
     }),
     'context': <ANY>,
@@ -420,13 +421,13 @@
       'name': None,
       'options': dict({
       }),
-      'original_device_class': None,
+      'original_device_class': <SwitchDeviceClass.SWITCH: 'switch'>,
       'original_icon': None,
       'original_name': None,
       'platform': 'vesync',
       'previous_unique_id': None,
       'supported_features': 0,
-      'translation_key': None,
+      'translation_key': 'on',
       'unique_id': 'switch',
       'unit_of_measurement': None,
     }),
@@ -435,6 +436,7 @@
 # name: test_switch_state[Wall Switch][switch.wall_switch]
   StateSnapshot({
     'attributes': ReadOnlyDict({
+      'device_class': 'switch',
       'friendly_name': 'Wall Switch',
     }),
     'context': <ANY>,