Rebase from upstream

custom_components/localtuya/common.py

@@ -13,6 +13,7 @@ from homeassistant.const import (
     CONF_ID,
     CONF_PLATFORM,
     CONF_SCAN_INTERVAL,
+    STATE_UNKNOWN,
 )
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import (
@@ -31,6 +32,9 @@ from .const import (
     DATA_CLOUD,
     DOMAIN,
     TUYA_DEVICES,
+    CONF_DEFAULT_VALUE,
+    ATTR_STATE,
+    CONF_RESTORE_ON_RECONNECT,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -91,6 +95,8 @@ async def async_setup_entry(
                         entity_config[CONF_ID],
                     )
                 )
+    #Once the entities have been created, add to the TuyaDevice instance
+    tuyainterface.add_entities(entities)
     async_add_entities(entities)
 
 
@@ -135,6 +141,7 @@ class TuyaDevice(pytuya.TuyaListener, pytuya.ContextualLogger):
         self._connect_task = None
         self._disconnect_task = None
         self._unsub_interval = None
+        self._entities = []
         self._local_key = self._dev_config_entry[CONF_LOCAL_KEY]
         self.set_logger(_LOGGER, self._dev_config_entry[CONF_DEVICE_ID])
 
@@ -142,6 +149,10 @@ class TuyaDevice(pytuya.TuyaListener, pytuya.ContextualLogger):
         for entity in self._dev_config_entry[CONF_ENTITIES]:
             self.dps_to_request[entity[CONF_ID]] = None
 
+    def add_entities(self, entities):
+        """Set the entities associated with this device"""
+        self._entities.extend(entities)
+
     @property
     def connected(self):
         """Return if connected to device."""
@@ -173,6 +184,10 @@ class TuyaDevice(pytuya.TuyaListener, pytuya.ContextualLogger):
 
             self.status_updated(status)
 
+            # Attempt to restore status for all entites that need to first set the DPS value before the device will respond with status.
+            for entity in self._entities:
+                await entity.restore_state_when_connected()
+
             def _new_entity_handler(entity_id):
                 self.debug(
                     "New entity %s was added to %s",
@@ -254,7 +269,7 @@ class TuyaDevice(pytuya.TuyaListener, pytuya.ContextualLogger):
             try:
                 await self._interface.set_dp(state, dp_index)
             except Exception:  # pylint: disable=broad-except
-                self.exception("Failed to set DP %d to %d", dp_index, state)
+                self.exception("Failed to set DP %d to %s", dp_index, str(state))
         else:
             self.error(
                 "Not connected to device %s", self._dev_config_entry[CONF_FRIENDLY_NAME]
@@ -305,6 +320,16 @@ class LocalTuyaEntity(RestoreEntity, pytuya.ContextualLogger):
         self._config = get_entity_config(config_entry, dp_id)
         self._dp_id = dp_id
         self._status = {}
+        self._state = None
+        self._last_state = None
+
+        #Default value is available to be provided by Platform entities if required
+        self._default_value = self._config.get(CONF_DEFAULT_VALUE)
+
+        #Restore on connect setting is available to be provided by Platform entities if required
+        self._restore_on_reconnect = (
+            self._config.get(CONF_RESTORE_ON_RECONNECT) or False
+        )
         self.set_logger(logger, self._dev_config_entry[CONF_DEVICE_ID])
 
     async def async_added_to_hass(self):
@@ -325,6 +350,8 @@ class LocalTuyaEntity(RestoreEntity, pytuya.ContextualLogger):
                 self._status = status.copy()
                 if status:
                     self.status_updated()
+
+                # Update HA
                 self.schedule_update_ha_state()
 
         signal = f"localtuya_{self._dev_config_entry[CONF_DEVICE_ID]}"
@@ -336,6 +363,20 @@ class LocalTuyaEntity(RestoreEntity, pytuya.ContextualLogger):
         signal = f"localtuya_entity_{self._dev_config_entry[CONF_DEVICE_ID]}"
         async_dispatcher_send(self.hass, signal, self.entity_id)
 
+    @property
+    def extra_state_attributes(self):
+        """Return entity specific state attributes to be saved & then available for restore
+           when the entity is restored at startup.
+        """
+        attributes = {}
+        if self._state is not None:
+            attributes[ATTR_STATE] = self._state
+        elif self._last_state is not None:
+            attributes[ATTR_STATE] = self._last_state
+
+        self.debug("Entity %s - Additional attributes: %s", self.name, attributes)
+        return attributes
+
     @property
     def device_info(self):
         """Return device information for the device registry."""
@@ -408,9 +449,83 @@ class LocalTuyaEntity(RestoreEntity, pytuya.ContextualLogger):
 
         Override in subclasses and update entity specific state.
         """
+        state = self.dps(self._dp_id)
+        self._state = state
+
+        # Keep record in last_state as long as not during connection/re-connection,
+        # as last state will be used to restore the previous state
+        if (state is not None) and (self._device._connect_task is None):
+            self._last_state = state
 
     def status_restored(self, stored_state):
         """Device status was restored.
 
         Override in subclasses and update entity specific state.
         """
+        raw_state = stored_state.attributes.get(ATTR_STATE)
+        if raw_state is not None:
+            #    (stored_state.state == "unavailable") | (stored_state.state == "unknown")
+            # ):
+            self._last_state = raw_state
+            self.debug(
+                "Restoring state for entity: %s - state: %s",
+                self.name,
+                str(self._last_state),
+            )
+
+    def default_value(self):
+        """Default value of this entity
+
+        Override in subclasses to specify the default value for the entity.
+        """
+        # Check if default value has been set - if not, default to the entity defaults.
+        if self._default_value is None:
+            self._default_value = self.entity_default_value()
+
+        return self._default_value
+
+    def entity_default_value(self):
+        """Default value of the entity type
+
+        Override in subclasses to specify the default value for the entity.
+        """
+        return 0
+
+    @property
+    def restore_on_reconnect(self):
+        """Returns whether the last state should be restored on a reconnect - useful where the device loses settings if powered off"""
+        return self._restore_on_reconnect
+
+    async def restore_state_when_connected(self):
+        """Restore if restore_on_reconnect is set, or if no status has been yet found - which indicates a DPS that needs to be set before it starts returning status"""
+        if not self.restore_on_reconnect and (str(self._dp_id) in self._status):
+            self.debug(
+                "Entity %s (DP %d) - Not restoring as restore on reconnect is disabled for this entity and the entity has an initial status",
+                self.name,
+                self._dp_id,
+            )
+            return
+
+        self.debug("Attempting to restore state for entity: %s", self.name)
+        # Attempt to restore the current state - in case reset.
+        restore_state = self._state
+
+        # If no state stored in the entity currently, go from last saved state
+        if (restore_state == STATE_UNKNOWN) | (restore_state is None):
+            self.debug("No current state for entity")
+            restore_state = self._last_state
+
+        # If no current or saved state, then use the default value
+        if restore_state is None:
+            self.debug("No last restored state - using default")
+            restore_state = self.default_value()
+
+        self.debug(
+            "Entity %s (DP %d) - Restoring state: %s",
+            self.name,
+            self._dp_id,
+            str(restore_state),
+        )
+
+        # Manually initialise
+        await self._device.set_dp(restore_state, self._dp_id)