From 1b73cfce3a3faea3eb6dc529f44814f0e2ffeaa8 Mon Sep 17 00:00:00 2001
From: Felipe Martins Diel <41558831+felipediel@users.noreply.github.com>
Date: Tue, 16 Feb 2021 16:38:10 -0300
Subject: [PATCH] Split the sp2 class into smaller classes (#521)

---
 broadlink/__init__.py | 19 +++++++++----------
 broadlink/switch.py   | 44 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/broadlink/__init__.py b/broadlink/__init__.py
index 1b2f0ab..0713127 100644
--- a/broadlink/__init__.py
+++ b/broadlink/__init__.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """The python-broadlink library."""
 import socket
-from typing import Generator, List, Union, Tuple
+from typing import Generator, List, Tuple, Union
 
 from .alarm import S1C
 from .climate import hysen
@@ -11,21 +11,16 @@ from .exceptions import exception
 from .light import lb1
 from .remote import rm, rm4
 from .sensor import a1
-from .switch import bg1, mp1, sp1, sp2, sp4, sp4b
-
+from .switch import bg1, mp1, sp1, sp2, sp2s, sp3, sp3s, sp4, sp4b
 
 SUPPORTED_TYPES = {
     0x0000: (sp1, "SP1", "Broadlink"),
     0x2711: (sp2, "SP2", "Broadlink"),
-    0x2716: (sp2, "NEO PRO", "Ankuoo"),
     0x2717: (sp2, "NEO", "Ankuoo"),
     0x2719: (sp2, "SP2-compatible", "Honeywell"),
     0x271A: (sp2, "SP2-compatible", "Honeywell"),
-    0x271D: (sp2, "Ego", "Efergy"),
     0x2720: (sp2, "SP mini", "Broadlink"),
     0x2728: (sp2, "SP2-compatible", "URANT"),
-    0x2733: (sp2, "SP3", "Broadlink"),
-    0x2736: (sp2, "SP mini+", "Broadlink"),
     0x273E: (sp2, "SP mini", "Broadlink"),
     0x7530: (sp2, "SP2", "Broadlink (OEM)"),
     0x7539: (sp2, "SP2-IL", "Broadlink (OEM)"),
@@ -37,10 +32,14 @@ SUPPORTED_TYPES = {
     0x7918: (sp2, "SP2", "Broadlink (OEM)"),
     0x7919: (sp2, "SP2-compatible", "Honeywell"),
     0x791A: (sp2, "SP2-compatible", "Honeywell"),
-    0x7D00: (sp2, "SP3-EU", "Broadlink (OEM)"),
     0x7D0D: (sp2, "SP mini 3", "Broadlink (OEM)"),
-    0x9479: (sp2, "SP3S-US", "Broadlink"),
-    0x947A: (sp2, "SP3S-EU", "Broadlink"),
+    0x2716: (sp2s, "NEO PRO", "Ankuoo"),
+    0x271D: (sp2s, "Ego", "Efergy"),
+    0x2736: (sp2s, "SP mini+", "Broadlink"),
+    0x2733: (sp3, "SP3", "Broadlink"),
+    0x7D00: (sp3, "SP3-EU", "Broadlink (OEM)"),
+    0x9479: (sp3s, "SP3S-US", "Broadlink"),
+    0x947A: (sp3s, "SP3S-EU", "Broadlink"),
     0x756C: (sp4, "SP4M", "Broadlink"),
     0x756F: (sp4, "MCB1", "Broadlink"),
     0x7579: (sp4, "SP4L-EU", "Broadlink"),
diff --git a/broadlink/switch.py b/broadlink/switch.py
index e5e0a1a..775b758 100644
--- a/broadlink/switch.py
+++ b/broadlink/switch.py
@@ -160,6 +160,44 @@ class sp2(device):
 
     TYPE = "SP2"
 
+    def set_power(self, state: bool) -> None:
+        """Set the power state of the device."""
+        packet = bytearray(16)
+        packet[0] = 2
+        packet[4] = int(bool(state))
+        response = self.send_packet(0x6A, packet)
+        check_error(response[0x22:0x24])
+
+    def check_power(self) -> bool:
+        """Return the power state of the device."""
+        packet = bytearray(16)
+        packet[0] = 1
+        response = self.send_packet(0x6A, packet)
+        check_error(response[0x22:0x24])
+        payload = self.decrypt(response[0x38:])
+        return bool(payload[0x4])
+
+
+class sp2s(sp2):
+    """Controls a Broadlink SP2S."""
+
+    TYPE = "SP2S"
+
+    def get_energy(self) -> float:
+        """Return the power consumption in W."""
+        packet = bytearray(16)
+        packet[0] = 4
+        response = self.send_packet(0x6A, packet)
+        check_error(response[0x22:0x24])
+        payload = self.decrypt(response[0x38:])
+        return int.from_bytes(payload[0x4:0x7], "little") / 1000
+
+
+class sp3(device):
+    """Controls a Broadlink SP3."""
+
+    TYPE = "SP3"
+
     def set_power(self, state: bool) -> None:
         """Set the power state of the device."""
         packet = bytearray(16)
@@ -200,6 +238,12 @@ class sp2(device):
         payload = self.decrypt(response[0x38:])
         return bool(payload[0x4] == 2 or payload[0x4] == 3 or payload[0x4] == 0xFF)
 
+
+class sp3s(sp2):
+    """Controls a Broadlink SP3S."""
+
+    TYPE = "SP3S"
+
     def get_energy(self) -> float:
         """Return the power consumption in W."""
         packet = bytearray([8, 0, 254, 1, 5, 1, 0, 0, 0, 45])