Switch cleanup for Shelly (part 1) (#138791)

homeassistant/components/shelly/switch.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any, cast
 
@@ -29,7 +30,7 @@ from .entity import (
     ShellyRpcEntity,
     ShellySleepingBlockAttributeEntity,
     async_setup_entry_attribute_entities,
-    async_setup_rpc_attribute_entities,
+    async_setup_entry_rpc,
 )
 from .utils import (
     async_remove_orphaned_entities,
@@ -60,18 +61,32 @@ MOTION_SWITCH = BlockSwitchDescription(
 class RpcSwitchDescription(RpcEntityDescription, SwitchEntityDescription):
     """Class to describe a RPC virtual switch."""
 
+    is_on: Callable[[dict[str, Any]], bool]
+    method_on: str
+    method_off: str
+    method_params_fn: Callable[[int | None, bool], dict]
 
-RPC_VIRTUAL_SWITCH = RpcSwitchDescription(
-    key="boolean",
-    sub_key="value",
-)
 
-RPC_SCRIPT_SWITCH = RpcSwitchDescription(
-    key="script",
-    sub_key="running",
-    entity_registry_enabled_default=False,
-    entity_category=EntityCategory.CONFIG,
-)
+RPC_SWITCHES = {
+    "boolean": RpcSwitchDescription(
+        key="boolean",
+        sub_key="value",
+        is_on=lambda status: bool(status["value"]),
+        method_on="Boolean.Set",
+        method_off="Boolean.Set",
+        method_params_fn=lambda id, value: {"id": id, "value": value},
+    ),
+    "script": RpcSwitchDescription(
+        key="script",
+        sub_key="running",
+        is_on=lambda status: bool(status["running"]),
+        method_on="Script.Start",
+        method_off="Script.Stop",
+        method_params_fn=lambda id, _: {"id": id},
+        entity_registry_enabled_default=False,
+        entity_category=EntityCategory.CONFIG,
+    ),
+}
 
 
 async def async_setup_entry(
@@ -174,20 +189,8 @@ def async_setup_rpc_entry(
         unique_id = f"{coordinator.mac}-switch:{id_}"
         async_remove_shelly_entity(hass, "light", unique_id)
 
-    async_setup_rpc_attribute_entities(
-        hass,
-        config_entry,
-        async_add_entities,
-        {"boolean": RPC_VIRTUAL_SWITCH},
-        RpcVirtualSwitch,
-    )
-
-    async_setup_rpc_attribute_entities(
-        hass,
-        config_entry,
-        async_add_entities,
-        {"script": RPC_SCRIPT_SWITCH},
-        RpcScriptSwitch,
+    async_setup_entry_rpc(
+        hass, config_entry, async_add_entities, RPC_SWITCHES, RpcSwitch
     )
 
     # the user can remove virtual components from the device configuration, so we need
@@ -324,8 +327,8 @@ class RpcRelaySwitch(ShellyRpcEntity, SwitchEntity):
         await self.call_rpc("Switch.Set", {"id": self._id, "on": False})
 
 
-class RpcVirtualSwitch(ShellyRpcAttributeEntity, SwitchEntity):
-    """Entity that controls a virtual boolean component on RPC based Shelly devices."""
+class RpcSwitch(ShellyRpcAttributeEntity, SwitchEntity):
+    """Entity that controls a switch on RPC based Shelly devices."""
 
     entity_description: RpcSwitchDescription
     _attr_has_entity_name = True
@@ -333,32 +336,18 @@ class RpcVirtualSwitch(ShellyRpcAttributeEntity, SwitchEntity):
     @property
     def is_on(self) -> bool:
         """If switch is on."""
-        return bool(self.attribute_value)
+        return self.entity_description.is_on(self.status)
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn on relay."""
-        await self.call_rpc("Boolean.Set", {"id": self._id, "value": True})
+        await self.call_rpc(
+            self.entity_description.method_on,
+            self.entity_description.method_params_fn(self._id, True),
+        )
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off relay."""
-        await self.call_rpc("Boolean.Set", {"id": self._id, "value": False})
-
-
-class RpcScriptSwitch(ShellyRpcAttributeEntity, SwitchEntity):
-    """Entity that controls a script component on RPC based Shelly devices."""
-
-    entity_description: RpcSwitchDescription
-    _attr_has_entity_name = True
-
-    @property
-    def is_on(self) -> bool:
-        """If switch is on."""
-        return bool(self.status["running"])
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn on relay."""
-        await self.call_rpc("Script.Start", {"id": self._id})
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn off relay."""
-        await self.call_rpc("Script.Stop", {"id": self._id})
+        await self.call_rpc(
+            self.entity_description.method_off,
+            self.entity_description.method_params_fn(self._id, False),
+        )