Persist pyscript.* states between home assistant restarts

With this change any states set by pyscript in the pyscript domain
will be persisted when home assistant next reboots.

For example when a script runs `pyscript.abc = 'foo'` then this value
will be persisted.

Fixes #47

Author: swazrgb

custom_components/pyscript/__init__.py

@@ -17,6 +17,7 @@
 )
 from homeassistant.exceptions import HomeAssistantError
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.restore_state import RestoreStateData
 from homeassistant.loader import bind_hass
 
 from .const import CONF_ALLOW_ALL_IMPORTS, DOMAIN, FOLDER, LOGGER_PATH, SERVICE_JUPYTER_KERNEL_START
@@ -39,6 +40,7 @@
 
 async def async_setup(hass, config):
     """Component setup, run import config flow for each entry in config."""
+    await restore_state(hass)
     if DOMAIN in config:
         hass.async_create_task(
             hass.config_entries.flow.async_init(
@@ -49,6 +51,14 @@ async def async_setup(hass, config):
     return True
 
 
+async def restore_state(hass):
+    restore_data = await RestoreStateData.async_get_instance(hass)
+    for entity_id, value in restore_data.last_states.items():
+        if entity_id.startswith("pyscript."):
+            last_state = value.state
+            hass.states.async_set(entity_id, last_state.state, last_state.attributes)
+
+
 async def async_setup_entry(hass, config_entry):
     """Initialize the pyscript config entry."""
     Function.init(hass)

custom_components/pyscript/eval.py

@@ -1168,7 +1168,7 @@ async def recurse_assign(self, lhs, val):
             if not isinstance(var_name, str):
                 raise NotImplementedError(f"unknown lhs type {lhs} (got {var_name}) in assign")
             if var_name.find(".") >= 0:
-                State.set(var_name, val)
+                await State.set(var_name, val)
                 return
             if self.curr_func and var_name in self.curr_func.global_names:
                 self.global_sym_table[var_name] = val

custom_components/pyscript/state.py

@@ -2,6 +2,8 @@
 
 import logging
 
+from homeassistant.helpers.restore_state import RestoreStateData
+
 from .const import LOGGER_PATH
 from .function import Function
 
@@ -98,7 +100,7 @@ def notify_var_get(cls, var_names, new_vars):
         return notify_vars
 
     @classmethod
-    def set(cls, var_name, value, new_attributes=None, **kwargs):
+    async def set(cls, var_name, value, new_attributes=None, **kwargs):
         """Set a state variable and optional attributes in hass."""
         if var_name.count(".") != 1:
             raise NameError(f"invalid name {var_name} (should be 'domain.entity')")
@@ -113,6 +115,12 @@ def set(cls, var_name, value, new_attributes=None, **kwargs):
             new_attributes.update(kwargs)
         _LOGGER.debug("setting %s = %s, attr = %s", var_name, value, new_attributes)
         cls.notify_var_last[var_name] = str(value)
+
+        if var_name.startswith("pyscript."):
+            # have this var tracked for restore
+            restore_data = await RestoreStateData.async_get_instance(cls.hass)
+            restore_data.async_restore_entity_added(var_name)
+
         cls.hass.states.async_set(var_name, value, new_attributes)
 
     @classmethod