Update spec for Neo Power Plug (zigbee)

custom_components/xiaomi_gateway3/core/converters/devices.py

@@ -917,7 +917,7 @@
 }, {
     # Neo Power Plug NAS-WR01B
     "TS011F": ["Neo", "Power Plug", "NAS-WR01B"],
-    "support": 3,
+    "support": 5,
     "spec": [
         ZOnOffConv("plug", "switch"),
         # default gateway software binds electrical_measurement and
@@ -926,8 +926,10 @@
         ZCurrentConv("current", "sensor", report="5s 1h 1", multiply=0.001),
         ZPowerConv("power", "sensor", report="5s 1h 1"),
         ZEnergyConv("energy", "sensor", report="5s 1h 1", multiply=0.01),
-        ZTuyaPowerOn,
-        ZTuyaModeConv("led", "select", enabled=False)
+        ZTuyaPowerOnConv("power_on_state", "select", enabled=False),
+        ZTuyaLEDModeConv("led", "select", enabled=False),
+        ZTuyaChildModeConv("child_mode", "switch", enabled=False),
+        ZTuyaModeConv("mode", "select", enabled=False)
     ],
 }, {
     # tuya relay with neutral, 1 gang

custom_components/xiaomi_gateway3/core/converters/silabs.py

@@ -282,14 +282,20 @@ def zcl_color(nwk: str, ep: int, ct: int, tr: float) -> list:
 
 
 # zcl global read [cluster:2] [attributeId:2]
-def zcl_read(nwk: str, ep: int, cluster: str, *attrs) -> list:
+def zcl_read(nwk: str, ep: int, cluster: Union[str, int], *attrs) -> list:
     """Generate Silabs Z3 read attribute command. Support multiple attrs."""
-    # convert string to object
-    cluster = get_cluster(cluster)
-    cid = cluster.cluster_id
+    if isinstance(cluster, str):
+        cluster = get_cluster(cluster)
+        cid = cluster.cluster_id
+
+        # convert List[str] to List[int]
+        attrs = [get_attr(cluster.attributes, attr) for attr in attrs]
+    else:
+        cid = cluster
 
-    # convert List[str] to List[int]
-    attrs = [get_attr(cluster.attributes, attr) for attr in attrs]
+    assert isinstance(cid, int)
+    for attr in attrs:
+        assert isinstance(attr, int)
 
     if len(attrs) > 1:
         raw = "".join([int(a).to_bytes(2, "little").hex() for a in attrs])

custom_components/xiaomi_gateway3/core/converters/zigbee.py

@@ -85,6 +85,27 @@ def decode(self, device: 'XDevice', payload: dict, value: dict):
         if value["endpoint"] == self.ep and self.zattr in value:
             payload[self.attr] = bool(value[self.zattr])
 
+    def encode(self, device: "XDevice", payload: dict, value: bool):
+        cmd = zcl_write(
+            device.nwk, self.ep, self.zigbee, self.zattr, int(value), type=0x10
+        )
+        payload.setdefault("commands", []).extend(cmd)
+
+
+class ZMapConv(ZConverter):
+    map = {}
+
+    def decode(self, device: 'XDevice', payload: dict, value: dict):
+        if self.zattr in value:
+            payload[self.attr] = self.map.get(value[self.zattr])
+
+    def encode(self, device: "XDevice", payload: dict, value: str):
+        v = next(k for k, v in self.map.items() if v == value)
+        cmd = zcl_write(
+            device.nwk, self.ep, self.zigbee, self.zattr, v, type=0x30
+        )
+        payload.setdefault("commands", []).extend(cmd)
+
 
 @dataclass
 class ZMathConv(ZConverter):
@@ -250,30 +271,30 @@ def read(self, device: "XDevice", payload: dict):
 # Specific defices converters
 ################################################################################
 
+class ZTuyaChildModeConv(ZBoolConv):
+    zigbee = "on_off"
+    zattr = 0x8000
+
+
+class ZTuyaLEDModeConv(ZMapConv):
+    zigbee = "on_off"
+    zattr = 0x8001
+    map = {0: "off", 1: "off/on", 2: "on/off", 3: "on"}
+
+
 # Thanks to:
 # https://github.com/Koenkk/zigbee-herdsman/blob/master/src/zcl/definition/cluster.ts
 # moesStartUpOnOff: {ID: 0x8002, type: DataType.enum8},
-class ZTuyaPowerOnConv(ZConverter):
+class ZTuyaPowerOnConv(ZMapConv):
     zigbee = "on_off"
     zattr = 0x8002
     map = {0: "off", 1: "on", 2: "previous"}
 
-    def decode(self, device: 'XDevice', payload: dict, value: dict):
-        if self.zattr in value:
-            payload[self.attr] = self.map.get(value[self.zattr])
-
-    def encode(self, device: "XDevice", payload: dict, value: str):
-        v = next(k for k, v in self.map.items() if v == value)
-        cmd = zcl_write(
-            device.nwk, self.ep, self.zigbee, self.zattr, v, type=0x30
-        )
-        payload.setdefault("commands", []).extend(cmd)
-
 
 # Thanks to:
 # https://github.com/Koenkk/zigbee-herdsman-converters/blob/910271ae8fccb19305752d3f67381b4765853018/converters/fromZigbee.js#L4537
 # https://github.com/Koenkk/zigbee-herdsman/blob/068bbe7636f588394f69f82bc25c8b68a4feada7/src/zcl/definition/cluster.ts#L4284
-class ZTuyaModeConv(ZTuyaPowerOnConv):
+class ZTuyaModeConv(ZMapConv):
     zigbee = 0xE001
     zattr = 0xD030
     map = {0: "toggle", 1: "state", 2: "momentary"}

custom_components/xiaomi_gateway3/core/entity.py

@@ -80,6 +80,7 @@
     ZIGBEE: "mdi:zigbee",
     "action": "mdi:bell",
     "alarm": "mdi:shield-home",
+    "child_mode": "mdi:baby-carriage",
     "conductivity": "mdi:flower",
     "gas_density": "mdi:google-circles-communities",
     "group": "mdi:lightbulb-group",
@@ -122,14 +123,14 @@
     GATEWAY: ENTITY_CATEGORY_DIAGNOSTIC,
     MESH: ENTITY_CATEGORY_DIAGNOSTIC,
     ZIGBEE: ENTITY_CATEGORY_DIAGNOSTIC,
-    "baby_mode": ENTITY_CATEGORY_CONFIG,
     "battery": ENTITY_CATEGORY_DIAGNOSTIC,
     "battery_charging": ENTITY_CATEGORY_DIAGNOSTIC,
     "battery_low": ENTITY_CATEGORY_DIAGNOSTIC,
     "battery_percent": ENTITY_CATEGORY_DIAGNOSTIC,
     "battery_voltage": ENTITY_CATEGORY_DIAGNOSTIC,
     "blind_time": ENTITY_CATEGORY_CONFIG,
     "charge_protect": ENTITY_CATEGORY_CONFIG,
+    "child_mode": ENTITY_CATEGORY_CONFIG,
     "chip_temperature": ENTITY_CATEGORY_DIAGNOSTIC,
     "cloud_link": ENTITY_CATEGORY_DIAGNOSTIC,
     "display_unit": ENTITY_CATEGORY_CONFIG,