handle logbook

Author: dlashua

custom_components/pyscript/function.py

@@ -26,6 +26,10 @@ class Function:
     unique_task2name = {}
     unique_name2task = {}
 
+    #
+    # Mappings of task id to hass contexts
+    task2context = {}
+
     #
     # Set of tasks that are running
     #
@@ -117,12 +121,21 @@ async def async_sleep(cls, duration):
     @classmethod
     async def event_fire(cls, event_type, **kwargs):
         """Implement event.fire()."""
+        curr_task = asyncio.current_task()
         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)
+        elif curr_task in cls.task2context:
+            context = cls.task2context[curr_task]
         else:
-            cls.hass.bus.async_fire(event_type, kwargs)
+            context = Context()
+
+        cls.hass.bus.async_fire(event_type, kwargs, context=context)
+
+    @classmethod
+    def store_hass_context(cls, hass_context):
+        curr_task = asyncio.current_task()
+        cls.task2context[curr_task] = hass_context
 
     @classmethod
     def task_unique_factory(cls, ctx):
@@ -177,12 +190,16 @@ def service_has_service(cls, domain, name):
     @classmethod
     async def service_call(cls, domain, name, **kwargs):
         """Implement service.call()."""
+        curr_task = asyncio.current_task()
         if "context" in kwargs and isinstance(kwargs["context"], Context):
             context = kwargs["context"]
             del kwargs["context"]
-            await cls.hass.services.async_call(domain, name, kwargs, context=context)
+        elif curr_task in cls.task2context:
+            context = cls.task2context[curr_task]
         else:
-            await cls.hass.services.async_call(domain, name, kwargs)
+            context = Context()
+
+        await cls.hass.services.async_call(domain, name, kwargs, context=context)
 
     @classmethod
     async def service_completions(cls, root):
@@ -239,12 +256,16 @@ def get(cls, name):
             return None
 
         async def service_call(*args, **kwargs):
+            curr_task = asyncio.current_task()
             if "context" in kwargs and isinstance(kwargs["context"], Context):
                 context = kwargs["context"]
                 del kwargs["context"]
-                await cls.hass.services.async_call(domain, service, kwargs, False, context)
+            elif curr_task in cls.task2context:
+                context = cls.task2context[curr_task]
             else:
-                await cls.hass.services.async_call(domain, service, kwargs)
+                context = Context()
+
+            await cls.hass.services.async_call(domain, service, kwargs, context=context)
 
         return service_call
 
@@ -266,6 +287,8 @@ async def run_coro(cls, coro):
             if task in cls.unique_task2name:
                 del cls.unique_name2task[cls.unique_task2name[task]]
                 del cls.unique_task2name[task]
+            if task in cls.task2context:
+                del cls.task2context[task]
             cls.our_tasks.discard(task)
 
     @classmethod

custom_components/pyscript/logbook.py

@@ -0,0 +1,45 @@
+"""Describe logbook events."""
+from homeassistant.core import callback
+
+from .const import DOMAIN
+
+import logging
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@callback
+def async_describe_events(hass, async_describe_event):  # type: ignore
+    """Describe logbook events."""
+
+    @callback
+    def async_describe_logbook_event(event):  # type: ignore
+        """Describe a logbook event."""
+        data = event.data
+        func_args = data.get("func_args", {})
+        ev_name = data.get("name", "unknown")
+        ev_entity_id = data.get("entity_id", "pyscript.unknown")
+
+        ev_trigger_type = func_args.get("trigger_type", "unknown")
+        if ev_trigger_type == "event":
+            ev_source = f"event {func_args.get('event_type', 'unknown event')}"
+        elif ev_trigger_type == "state":
+            ev_source = f"state change {func_args.get('var_name', 'unknown entity')} == {func_args.get('value', 'unknown value')}"
+        elif ev_trigger_type == "time":
+            ev_trigger_time = func_args.get("trigger_time", "unknown")
+            if ev_trigger_time is None:
+                ev_trigger_time = "startup"
+            ev_source = f"time {ev_trigger_time}"
+        else:
+            ev_source = ev_trigger_type
+
+        message = f"has been triggered by {ev_source}"
+
+        return {
+            "name": ev_name,
+            "message": message,
+            "source": ev_source,
+            "entity_id": ev_entity_id,
+        }
+
+    async_describe_event(DOMAIN, "pyscript_running", async_describe_logbook_event)

custom_components/pyscript/state.py

@@ -1,5 +1,6 @@
 """Handles state variable access and change notification."""
 
+import asyncio
 import logging
 
 from homeassistant.core import Context
@@ -117,10 +118,14 @@ async def set(cls, var_name, value, new_attributes=None, **kwargs):
             else:
                 new_attributes = {}
 
-        context = None
+        curr_task = asyncio.current_task()
         if "context" in kwargs and isinstance(kwargs["context"], Context):
             context = kwargs["context"]
             del kwargs["context"]
+        elif curr_task in Function.task2context:
+            context = Function.task2context[curr_task]
+        else:
+            context = Context()
 
         if kwargs:
             new_attributes = new_attributes.copy()

custom_components/pyscript/trigger.py

@@ -11,6 +11,7 @@
 from croniter import croniter
 
 import homeassistant.helpers.sun as sun
+from homeassistant.core import Context
 
 from .const import LOGGER_PATH
 from .eval import AstEval
@@ -826,14 +827,31 @@ async def trigger_watch(self):
                     )
                     continue
 
+                # Create new HASS Context with incoming as parent
+                if "context" in func_args and isinstance(func_args["context"], Context):
+                    hass_context = Context(parent_id=func_args["context"].id)
+                else:
+                    hass_context = Context()
+
+                # Fire an event indicating that pyscript is running
+                # Note: the event must have an entity_id for logbook to work correctly.
+                ev_name = self.name.replace(".", "_")
+                ev_entity_id = f"pyscript.{ev_name}"
+
+                event_data = dict(name=ev_name, entity_id=ev_entity_id, func_args=func_args)
+                Function.hass.bus.async_fire("pyscript_running", event_data, context=hass_context)
+
                 _LOGGER.debug(
                     "trigger %s got %s trigger, running action (kwargs = %s)",
                     self.name,
                     notify_type,
                     func_args,
                 )
 
-                async def do_func_call(func, ast_ctx, task_unique, task_unique_func, **kwargs):
+                async def do_func_call(func, ast_ctx, task_unique, task_unique_func, hass_context, **kwargs):
+                    # Store HASS Context for this Task
+                    Function.store_hass_context(hass_context)
+                    
                     if task_unique and task_unique_func:
                         await task_unique_func(task_unique)
                     await func.call(ast_ctx, **kwargs)
@@ -844,7 +862,12 @@ async def do_func_call(func, ast_ctx, task_unique, task_unique_func, **kwargs):
 
                 Function.create_task(
                     do_func_call(
-                        self.action, action_ast_ctx, self.task_unique, task_unique_func, **func_args
+                        self.action,
+                        action_ast_ctx,
+                        self.task_unique,
+                        task_unique_func,
+                        hass_context,
+                        **func_args,
                     )
                 )