event.fire and state.set send context if provided

Author: dlashua

custom_components/pyscript/function.py

@@ -117,7 +117,12 @@ async def async_sleep(cls, duration):
     @classmethod
     async def event_fire(cls, event_type, **kwargs):
         """Implement event.fire()."""
-        cls.hass.bus.async_fire(event_type, kwargs)
+        if "context" in kwargs and isinstance(kwargs["context"], Context):
+            context = kwargs["context"]
+            del kwargs["context"]
+            cls.hass.bus.async_fire(event_type, kwargs, context=context)
+        else:
+            cls.hass.bus.async_fire(event_type, kwargs)
 
     @classmethod
     def task_unique_factory(cls, ctx):

custom_components/pyscript/state.py

@@ -3,6 +3,7 @@
 import logging
 
 from homeassistant.helpers.restore_state import RestoreStateData
+from homeassistant.core import Context
 
 from .const import LOGGER_PATH
 from .function import Function
@@ -115,13 +116,19 @@ async def set(cls, var_name, value, new_attributes=None, **kwargs):
                 new_attributes = state_value.attributes
             else:
                 new_attributes = {}
+
+        context = None
+        if "context" in kwargs and isinstance(kwargs["context"], Context):
+            context = kwargs["context"]
+            del kwargs["context"]
+
         if kwargs:
             new_attributes = new_attributes.copy()
             new_attributes.update(kwargs)
         _LOGGER.debug("setting %s = %s, attr = %s", var_name, value, new_attributes)
         cls.notify_var_last[var_name] = str(value)
 
-        cls.hass.states.async_set(var_name, value, new_attributes)
+        cls.hass.states.async_set(var_name, value, new_attributes, context=context)
 
     @classmethod
     async def register_persist(cls, var_name):