Added restoring state after restart of HA

README.md

@@ -33,13 +33,16 @@ saver:
 ## Available services
 
 ### Save state
+Saves the state and parameters of the entity.
 ```yaml
 service: saver.save_state
 data:
   entity_id: cover.living_room
 ```
 
 ### Restore state
+Executes the script using saved state of the entity.
+
 To use state of an entity you have to use a following value in a template: `state`.
 
 To use state attribute (in this case `current_position`) of an entity you have to use a following value in a template: `attr_current_position`.
@@ -60,27 +63,31 @@ data:
 ```
 
 ### Delete saved state
+Deletes a saved state for an entity.
 ```yaml
 service: saver.delete
 data:
   entity_id: cover.living_room
 ```
 
 ### Set variable
+Sets the value to the variable.
 ```yaml
 service: saver.set_variable
 data:
   name: counter
 ```
 
 ### Delete variable
+Deletes a saved variable.
 ```yaml
 service: saver.delete_variable
 data:
   name: counter
 ```
 
 ### Execute script
+Executes a script using all saved entities and variables.
 
 To use state of an entity (in this case `cover.living_room`) you have to use a following value in a template: `cover_living_room_state`.
 
@@ -104,5 +111,10 @@ data:
         value: "Counter has value {{ counter }}"
 ```
 
+### Clear
+Deletes all saved data.
+```yaml
+service: saver.clear
+```
 
 <a href="https://www.buymeacoffee.com/PiotrMachowski" target="_blank"><img src="https://bmc-cdn.nyc3.digitaloceanspaces.com/BMC-button-images/custom_images/orange_img.png" alt="Buy Me A Coffee" style="height: auto !important;width: auto !important;" ></a>

custom_components/saver/__init__.py

@@ -1,7 +1,8 @@
 import logging
+import asyncio
 
 from homeassistant.core import Context
-from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.script import Script
 from homeassistant.helpers.entity_component import EntityComponent
 
@@ -26,6 +27,9 @@ def setup_entry(hass, config_entry):
     saver_entity = SaverEntity()
     component.add_entities([saver_entity])
 
+    def clear(call):
+        saver_entity.clear()
+
     def delete(call):
         data = call.data
         entity_id = data[CONF_ENTITY_ID]
@@ -59,6 +63,7 @@ def set_variable(call):
         value = data[CONF_VALUE]
         saver_entity.set_variable(name, value)
 
+    hass.services.register(DOMAIN, SERVICE_CLEAR, clear, SERVICE_CLEAR_SCHEMA)
     hass.services.register(DOMAIN, SERVICE_DELETE, delete, SERVICE_DELETE_SCHEMA)
     hass.services.register(DOMAIN, SERVICE_DELETE_VARIABLE, delete_variable, SERVICE_DELETE_VARIABLE_SCHEMA)
     hass.services.register(DOMAIN, SERVICE_EXECUTE, execute, SERVICE_EXECUTE_SCHEMA)
@@ -69,7 +74,7 @@ def set_variable(call):
     return True
 
 
-class SaverEntity(Entity):
+class SaverEntity(RestoreEntity):
 
     def __init__(self):
         self._entities_db = {}
@@ -79,6 +84,11 @@ def __init__(self):
     def name(self):
         return DOMAIN
 
+    def clear(self):
+        self._entities_db.clear()
+        self._variables_db.clear()
+        self.schedule_update_ha_state()
+
     def delete(self, entity_id):
         self._entities_db.pop(entity_id)
         self.schedule_update_ha_state()
@@ -118,20 +128,32 @@ def set_variable(self, variable, value):
     @property
     def state_attributes(self):
         return {
-            "entities": list(self._entities_db.keys()),
-            "variables": list(self._variables_db.keys())
+            "entities": self._entities_db,
+            "variables": self._variables_db
         }
 
     @property
     def state(self):
         return len(self._entities_db) + len(self._variables_db)
 
+    @asyncio.coroutine
+    def async_added_to_hass(self):
+        state = yield from self.async_get_last_state()
+        if state is not None \
+                and state.attributes is not None \
+                and "variables" in state.attributes and not isinstance(state.attributes["entities"], list) \
+                and "entities" in state.attributes and not isinstance(state.attributes["variables"], list):
+            self._variables_db = state.attributes["variables"]
+            self._entities_db = state.attributes["entities"]
+
     @staticmethod
     def convert_to_variables(state, entity_id=None):
         prefix = ""
+        state_val = state["state"] if isinstance(state, dict) else state.state
+        attrs = state["attributes"] if isinstance(state, dict) else state.attributes
         if entity_id is not None:
             prefix = f"{entity_id}_".replace(".", "_")
-        variables = {f"{prefix}state": state.state}
-        for attr in state.attributes:
-            variables[f"{prefix}attr_{attr}"] = state.attributes[attr]
+        variables = {f"{prefix}state": state_val}
+        for attr in attrs:
+            variables[f"{prefix}attr_{attr}"] = attrs[attr]
         return variables

custom_components/saver/const.py

@@ -10,6 +10,10 @@
 CONF_SCRIPT = 'script'
 CONF_VALUE = 'value'
 
+SERVICE_CLEAR = 'clear'
+SERVICE_CLEAR_SCHEMA = vol.Schema({
+})
+
 SERVICE_DELETE = 'delete'
 SERVICE_DELETE_SCHEMA = vol.Schema({
     vol.Required(CONF_ENTITY_ID): cv.entity_id

custom_components/saver/services.yaml

@@ -1,3 +1,7 @@
+clear:
+  description: >
+    Deletes all saved data.
+
 delete:
   description: >
     Deletes a saved state for an entity.