Make fan platform inherit from LocalTuyaEntity

Also big clean up and some bug fixes.

custom_components/localtuya/fan.py

@@ -15,25 +15,28 @@ fan:
 
 """
 import logging
-
+from time import time, sleep
-import voluptuous as vol
+from threading import Lock
 
 from homeassistant.components.fan import (
     FanEntity,
     DOMAIN,
     PLATFORM_SCHEMA,
+    SPEED_OFF,
     SPEED_LOW,
     SPEED_MEDIUM,
     SPEED_HIGH,
     SUPPORT_SET_SPEED,
     SUPPORT_OSCILLATE,
-    SUPPORT_DIRECTION,
 )
-from homeassistant.const import CONF_ID, CONF_FRIENDLY_NAME, STATE_OFF
+from homeassistant.const import CONF_ID, CONF_FRIENDLY_NAME
-import homeassistant.helpers.config_validation as cv
 
-from . import BASE_PLATFORM_SCHEMA, prepare_setup_entities, import_from_yaml
+from . import (
-from .pytuya import TuyaDevice
+    BASE_PLATFORM_SCHEMA,
+    LocalTuyaEntity,
+    prepare_setup_entities,
+    import_from_yaml,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -47,9 +50,7 @@ def flow_schema(dps):
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Setup a Tuya fan based on a config entry."""
-    device, entities_to_setup = prepare_setup_entities(
+    device, entities_to_setup = prepare_setup_entities(config_entry, DOMAIN)
-        config_entry, DOMAIN
-    )
     if not entities_to_setup:
         return
 
@@ -58,8 +59,8 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     for device_config in entities_to_setup:
         fans.append(
             LocaltuyaFan(
-                device,
+                TuyaCache(device, config_entry.data[CONF_FRIENDLY_NAME]),
-                device_config[CONF_FRIENDLY_NAME],
+                config_entry,
                 device_config[CONF_ID],
             )
         )
@@ -72,39 +73,101 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
     return import_from_yaml(hass, config, DOMAIN)
 
 
-class LocaltuyaFan(FanEntity):
+# TODO: Needed only to be compatible with LocalTuyaEntity
+class TuyaCache:
+    """Cache wrapper for pytuya.TuyaDevice"""
+
+    def __init__(self, device, friendly_name):
+        """Initialize the cache."""
+        self._cached_status = ""
+        self._cached_status_time = 0
+        self._device = device
+        self._friendly_name = friendly_name
+        self._lock = Lock()
+
+    @property
+    def unique_id(self):
+        """Return unique device identifier."""
+        return self._device.id
+
+    def __get_status(self):
+        for i in range(5):
+            try:
+                status = self._device.status()
+                return status
+            except Exception:
+                print(
+                    "Failed to update status of device [{}]".format(
+                        self._device.address
+                    )
+                )
+                sleep(1.0)
+                if i + 1 == 3:
+                    _LOGGER.error(
+                        "Failed to update status of device %s", self._device.address
+                    )
+                    #                    return None
+                    raise ConnectionError("Failed to update status .")
+
+    def set_dps(self, state, dps_index):
+        """Change the Tuya switch status and clear the cache."""
+        self._cached_status = ""
+        self._cached_status_time = 0
+        for i in range(5):
+            try:
+                return self._device.set_dps(state, dps_index)
+            except Exception:
+                print(
+                    "Failed to set status of device [{}]".format(self._device.address)
+                )
+                if i + 1 == 3:
+                    _LOGGER.error(
+                        "Failed to set status of device %s", self._device.address
+                    )
+                    return
+
+        #                    raise ConnectionError("Failed to set status.")
+        self.update()
+
+    def status(self):
+        """Get state of Tuya switch and cache the results."""
+        self._lock.acquire()
+        try:
+            now = time()
+            if not self._cached_status or now - self._cached_status_time > 15:
+                sleep(0.5)
+                self._cached_status = self.__get_status()
+                self._cached_status_time = time()
+            return self._cached_status
+        finally:
+            self._lock.release()
+
+
+class LocaltuyaFan(LocalTuyaEntity, FanEntity):
     """Representation of a Tuya fan."""
 
-    def __init__(self, device, friendly_name, switchid):
+    def __init__(
+        self,
+        device,
+        config_entry,
+        fanid,
+        **kwargs,
+    ):
         """Initialize the entity."""
-        self._device = device
+        super().__init__(device, config_entry, fanid, **kwargs)
-        self._name = friendly_name
+        self._is_on = False
-        self._available = False
+        self._speed = SPEED_OFF
-        self._friendly_name = friendly_name
-        self._switch_id = switchid
-        self._status = self._device.status()
-        self._state = False
-        self._supported_features = SUPPORT_SET_SPEED | SUPPORT_OSCILLATE
-        self._speed = STATE_OFF
         self._oscillating = False
 
     @property
     def oscillating(self):
         """Return current oscillating status."""
-        # if self._speed == STATE_OFF:
-        #    return False
-        _LOGGER.debug("localtuya fan: oscillating = %s", self._oscillating)
         return self._oscillating
 
-    @property
-    def name(self) -> str:
-        """Get entity name."""
-        return self._name
-
     @property
     def is_on(self):
-        """Check if Tuya switch is on."""
+        """Check if Tuya fan is on."""
-        return self._state
+        return self._is_on
 
     @property
     def speed(self) -> str:
@@ -114,96 +177,54 @@ class LocaltuyaFan(FanEntity):
     @property
     def speed_list(self) -> list:
         """Get the list of available speeds."""
-        return [STATE_OFF, SPEED_LOW, SPEED_MEDIUM, SPEED_HIGH]
+        return [SPEED_OFF, SPEED_LOW, SPEED_MEDIUM, SPEED_HIGH]
-        # return ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12']
 
     def turn_on(self, speed: str = None, **kwargs) -> None:
         """Turn on the entity."""
-        _LOGGER.debug("localtuya fan: turn_on speed to: %s", speed)
+        self._device.set_dps(True, "1")
-        # if speed is None:
-        #    speed = SPEED_MEDIUM
-        # self.set_speed(speed)
-        self._device.turn_on()
         if speed is not None:
             self.set_speed(speed)
-        self.schedule_update_ha_state()
+        else:
+            self.schedule_update_ha_state()
 
     def turn_off(self, **kwargs) -> None:
         """Turn off the entity."""
-        _LOGGER.debug("localtuya fan: turn_off")
+        self._device.set_dps(False, "1")
-        self._device.set_dps(False, '1')
-        self._state = False
         self.schedule_update_ha_state()
 
     def set_speed(self, speed: str) -> None:
         """Set the speed of the fan."""
-        _LOGGER.debug("localtuya fan: set_speed to: %s", speed)
         self._speed = speed
-        # self.schedule_update_ha_state()
+        if speed == SPEED_OFF:
-        if speed == STATE_OFF:
+            self._device.set_dps(False, "1")
-            self._device.set_dps(False, '1')
-            self._state = False
         elif speed == SPEED_LOW:
-            self._device.set_dps('1', '2')
+            self._device.set_dps("1", "2")
         elif speed == SPEED_MEDIUM:
-            self._device.set_dps('2', '2')
+            self._device.set_dps("2", "2")
         elif speed == SPEED_HIGH:
-            self._device.set_dps('3', '2')
+            self._device.set_dps("3", "2")
         self.schedule_update_ha_state()
 
-    # def set_direction(self, direction: str) -> None:
-    #    """Set the direction of the fan."""
-    #    self.direction = direction
-    #    self.schedule_update_ha_state()
-
     def oscillate(self, oscillating: bool) -> None:
         """Set oscillation."""
         self._oscillating = oscillating
         self._device.set_value("8", oscillating)
         self.schedule_update_ha_state()
 
-    # @property
-    # def current_direction(self) -> str:
-    #    """Fan direction."""
-    #    return self.direction
-
-    @property
-    def unique_id(self):
-        """Return unique device identifier."""
-        return f"local_{self._device.id}_{self._switch_id}"
-
-    @property
-    def available(self):
-        """Return if device is available or not."""
-        return self._available
-
     @property
     def supported_features(self) -> int:
         """Flag supported features."""
-        return self._supported_features
+        return SUPPORT_SET_SPEED | SUPPORT_OSCILLATE
 
-    def update(self):
+    def status_updated(self):
-        """Get state of Tuya switch."""
+        """Get state of Tuya fan."""
-        success = False
+        self._is_on = self._status["dps"]["1"]
-        for i in range(3):
+        if not self._status["dps"]["1"]:
-            if success is False:
+            self._speed = SPEED_OFF
-                try:
+        elif self._status["dps"]["2"] == "1":
-                    status = self._device.status()
+            self._speed = SPEED_LOW
-                    _LOGGER.debug("localtuya fan: status = %s", status)
+        elif self._status["dps"]["2"] == "2":
-                    self._state = status["dps"]["1"]
+            self._speed = SPEED_MEDIUM
-                    if status["dps"]["1"] == False:
+        elif self._status["dps"]["2"] == "3":
-                        self._speed = STATE_OFF
+            self._speed = SPEED_HIGH
-                    elif int(status["dps"]["2"]) == 1:
+        self._oscillating = self._status["dps"]["8"]
-                        self._speed = SPEED_LOW
-                    elif int(status["dps"]["2"]) == 2:
-                        self._speed = SPEED_MEDIUM
-                    elif int(status["dps"]["2"]) == 3:
-                        self._speed = SPEED_HIGH
-                    # self._speed = status['dps']['2']
-                    self._oscillating = status["dps"]["8"]
-                    success = True
-                except ConnectionError:
-                    if i + 1 == 3:
-                        success = False
-                        raise ConnectionError("Failed to update status.")
-        self._available = success