diff --git a/custom_components/pyscript/entity.py b/custom_components/pyscript/entity.py new file mode 100644 index 0000000..7b00e06 --- /dev/null +++ b/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 diff --git a/custom_components/pyscript/state.py b/custom_components/pyscript/state.py index 312ef95..12e2231 100644 --- a/custom_components/pyscript/state.py +++ b/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):