adds task.create(), task.cancel() and task.wait() per #112

Author: craigbarratt

custom_components/pyscript/eval.py

@@ -726,6 +726,10 @@ async def call(self, ast_ctx, *args, **kwargs):
         """Call the EvalFunc function."""
         return await self.func.call(ast_ctx, *args, **kwargs)
 
+    def get_name(self):
+        """Return the function name."""
+        return self.func.get_name()
+
     def __del__(self):
         """On deletion, stop any triggers for this function."""
         self.func.trigger_stop()
@@ -756,6 +760,10 @@ async def call(self, ast_ctx, *args, **kwargs):
         """Call the EvalFunc function."""
         return await self.eval_func_var.call(ast_ctx, *args, **kwargs)
 
+    def get_name(self):
+        """Return the function name."""
+        return self.eval_func_var.get_name()
+
     async def __call__(self, *args, **kwargs):
         """Call the EvalFunc function using our saved ast ctx."""
         return await self.eval_func_var.call(self.ast_ctx, *args, **kwargs)

custom_components/pyscript/function.py

@@ -63,11 +63,11 @@ def init(cls, hass):
         cls.hass = hass
         cls.functions.update(
             {
-                "task.executor": cls.task_executor,
                 "event.fire": cls.event_fire,
-                "task.sleep": cls.async_sleep,
                 "service.call": cls.service_call,
                 "service.has_service": cls.service_has_service,
+                "task.executor": cls.task_executor,
+                "task.sleep": cls.async_sleep,
             }
         )
         cls.ast_functions.update(
@@ -187,11 +187,7 @@ async def task_unique(name, kill_me=False):
                         await asyncio.sleep(100000)
                 elif task != curr_task and task in cls.our_tasks:
                     # only cancel tasks if they are ones we started
-                    try:
-                        task.cancel()
-                        await task
-                    except asyncio.CancelledError:
-                        pass
+                    Function.reaper_cancel(task)
             if curr_task in cls.our_tasks:
                 if name in cls.unique_name2task:
                     task = cls.unique_name2task[name]
@@ -325,7 +321,7 @@ async def run_coro(cls, coro):
         try:
             task = asyncio.current_task()
             cls.our_tasks.add(task)
-            await coro
+            return await coro
         except asyncio.CancelledError:
             raise
         except Exception:

custom_components/pyscript/trigger.py

@@ -14,7 +14,7 @@
 import homeassistant.helpers.sun as sun
 
 from .const import LOGGER_PATH
-from .eval import AstEval
+from .eval import AstEval, EvalFuncVar, EvalFuncVarAstCtx
 from .event import Event
 from .function import Function
 from .mqtt import Mqtt
@@ -133,11 +133,52 @@ async def wait_until_call(*arg, **kw):
 
             return wait_until_call
 
+        def user_task_create_factory(ast_ctx):
+            """Return wapper to call to astFunction with the ast context."""
+
+            async def user_task_create(func, *args, **kwargs):
+                """Implement task.create()."""
+
+                async def func_call(func, new_ast_ctx, *args, **kwargs):
+                    """Call user function inside task.create()."""
+                    ret = await func.call(new_ast_ctx, *args, **kwargs)
+                    if new_ast_ctx.get_exception_obj():
+                        new_ast_ctx.get_logger().error(new_ast_ctx.get_exception_long())
+                    return ret
+
+                if not isinstance(func, (EvalFuncVar, EvalFuncVarAstCtx)):
+                    raise TypeError("function is not callable by task.create()")
+
+                new_ast_ctx = AstEval(
+                    f"{ast_ctx.get_global_ctx_name()}.{func.get_name()}", ast_ctx.get_global_ctx()
+                )
+                Function.install_ast_funcs(new_ast_ctx)
+                return Function.create_task(func_call(func, new_ast_ctx, *args, **kwargs))
+
+            return user_task_create
+
         ast_funcs = {
             "task.wait_until": wait_until_factory,
+            "task.create": user_task_create_factory,
         }
         Function.register_ast(ast_funcs)
 
+        def user_task_cancel(task):
+            """Implement task.cancel()."""
+            if not isinstance(task, asyncio.Task):
+                raise TypeError(f"{task} is not of type asyncio.Task")
+            Function.reaper_cancel(task)
+
+        async def user_task_wait(aws):
+            """Implement task.wait()."""
+            return await asyncio.wait(aws)
+
+        funcs = {
+            "task.cancel": user_task_cancel,
+            "task.wait": user_task_wait,
+        }
+        Function.register(funcs)
+
         for i in range(0, 7):
             cls.dow2int[locale.nl_langinfo(getattr(locale, f"ABDAY_{i + 1}")).lower()] = i
             cls.dow2int[locale.nl_langinfo(getattr(locale, f"DAY_{i + 1}")).lower()] = i