From 9cb43dcc36238946c268335bf0c41cdafa332f22 Mon Sep 17 00:00:00 2001
From: Alone <hi@anlo.ng>
Date: Fri, 20 May 2022 16:34:57 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A2=20improve=20for=20number=20entity?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 custom_components/xiaomi_miot/number.py | 52 +++++++++++++++++--------
 1 file changed, 36 insertions(+), 16 deletions(-)

diff --git a/custom_components/xiaomi_miot/number.py b/custom_components/xiaomi_miot/number.py
index defd53204..5b4b1d7f9 100644
--- a/custom_components/xiaomi_miot/number.py
+++ b/custom_components/xiaomi_miot/number.py
@@ -6,6 +6,7 @@
     DOMAIN as ENTITY_DOMAIN,
     NumberEntity,
 )
+from homeassistant.helpers.restore_state import RestoreEntity, RestoreStateData
 
 from . import (
     DOMAIN,
@@ -54,32 +55,51 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
 class MiotNumberEntity(MiotEntity, NumberEntity):
     def __init__(self, config, miot_service: MiotService):
         super().__init__(miot_service, config=config, logger=_LOGGER)
+        self._attr_value = 0
 
-    @property
-    def value(self):
-        """Return the entity value to represent the entity state."""
-        return 0
-
-    def set_value(self, value: float):
+    def set_value(self, value):
         """Set new value."""
         raise NotImplementedError()
 
 
-class MiotNumberSubEntity(MiotPropertySubEntity, NumberEntity):
+class MiotNumberSubEntity(MiotPropertySubEntity, NumberEntity, RestoreEntity):
     def __init__(self, parent, miot_property: MiotProperty, option=None):
         super().__init__(parent, miot_property, option, domain=ENTITY_DOMAIN)
+        self._attr_value = 0
+        self._is_restore = False
+
+    async def async_added_to_hass(self):
+        await super().async_added_to_hass()
+        self._is_restore = self.custom_config_bool('restore_state')
+        if self._is_restore:
+            restored = await RestoreStateData.async_get_instance(self.hass)
+            if restored and self.entity_id in restored.last_states:
+                state = restored.last_states[self.entity_id].state
+                val = self.cast_value(state.state)
+                if val is not None:
+                    self._attr_value = val
+
+    def update(self, data=None):
+        super().update(data)
+        val = self.native_value
+        if val is not None:
+            self._attr_value = val
 
     @property
-    def value(self):
-        """Return the entity value to represent the entity state."""
+    def native_value(self):
         val = self._miot_property.from_dict(self._state_attrs)
+        return self.cast_value(val)
+
+    def cast_value(self, val, default=None):
         try:
-            val = float(val)
+            val = round(float(val), 6)
+            if self._miot_property.is_integer:
+                val = int(val)
         except (TypeError, ValueError):
-            val = 0
+            val = default
         return val
 
-    def set_value(self, value: float):
+    def set_value(self, value):
         """Set new value."""
         if self._miot_property.is_integer:
             value = int(value)
@@ -105,24 +125,24 @@ class MiotNumberActionSubEntity(MiotNumberSubEntity):
     def __init__(self, parent, miot_property: MiotProperty, miot_action: MiotAction, option=None):
         super().__init__(parent, miot_property, option)
         self._miot_action = miot_action
-        self._value = 0
         self._state_attrs.update({
             'miot_action': miot_action.full_name,
         })
 
     def update(self, data=None):
         self._available = True
-        self._value = 0
+        self._attr_value = 0
 
     @property
     def value(self):
         """Return the entity value to represent the entity state."""
-        return self._value
+        return self._attr_value
 
     def set_value(self, value: float):
         """Set new value."""
         val = int(value)
         ret = self.call_parent('call_action', self._miot_action, [val])
         if ret:
-            self._value = val
+            self._attr_value = val
+            self.async_write_ha_state()
         return ret