Merge pull request #332 from dlashua/fix-persist

fix for state.persist

Author: craigbarratt

custom_components/pyscript/entity.py

@@ -0,0 +1,19 @@
+"""Entity Classes"""
+from homeassistant.helpers.restore_state import RestoreEntity
+from homeassistant.helpers.typing import StateType
+from homeassistant.const import STATE_UNKNOWN
+
+
+class PyscriptEntity(RestoreEntity):
+    """Generic Pyscript Entity"""
+
+    _attr_extra_state_attributes: dict
+    _attr_state: StateType = STATE_UNKNOWN
+
+    def set_state(self, state):
+        """Set the state"""
+        self._attr_state = state
+
+    def set_attributes(self, attributes):
+        """Set Attributes"""
+        self._attr_extra_state_attributes = attributes

custom_components/pyscript/state.py

@@ -8,6 +8,7 @@
 from homeassistant.helpers.service import async_get_all_descriptions
 
 from .const import LOGGER_PATH
+from .entity import PyscriptEntity
 from .function import Function
 
 _LOGGER = logging.getLogger(LOGGER_PATH + ".state")
@@ -56,7 +57,7 @@ class State:
     #
     # pyscript vars which have already been registered as persisted
     #
-    persisted_vars = set()
+    persisted_vars = {}
 
     #
     # other parameters of all services that have "entity_id" as a parameter
@@ -198,6 +199,10 @@ def set(cls, var_name, value=None, new_attributes=None, **kwargs):
             #
             cls.notify_var_last[var_name] = StateVal(cls.hass.states.get(var_name))
 
+        if var_name in cls.persisted_vars:
+            cls.persisted_vars[var_name].set_state(value)
+            cls.persisted_vars[var_name].set_attributes(new_attributes)
+
     @classmethod
     def setattr(cls, var_attr_name, value):
         """Set a state variable's attribute in hass."""
@@ -213,8 +218,13 @@ async def register_persist(cls, var_name):
         """Register pyscript state variable to be persisted with RestoreState."""
         if var_name.startswith("pyscript.") and var_name not in cls.persisted_vars:
             restore_data = await RestoreStateData.async_get_instance(cls.hass)
-            restore_data.async_restore_entity_added(var_name)
-            cls.persisted_vars.add(var_name)
+            this_entity = PyscriptEntity()
+            this_entity.entity_id = var_name
+            cls.persisted_vars[var_name] = this_entity
+            try:
+                restore_data.async_restore_entity_added(this_entity)
+            except TypeError:
+                restore_data.async_restore_entity_added(var_name)
 
     @classmethod
     async def persist(cls, var_name, default_value=None, default_attributes=None):