added task.add_done_callback(), task.remove_done_callback() and tests; see #112

Author: craigbarratt

custom_components/pyscript/eval.py

@@ -762,6 +762,14 @@ def get_name(self):
         """Return the function name."""
         return self.eval_func_var.get_name()
 
+    def get_ast_ctx(self):
+        """Return the ast context."""
+        return self.ast_ctx
+
+    def get_eval_func_var(self):
+        """Return the eval_func_var."""
+        return self.eval_func_var
+
     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)
@@ -1703,25 +1711,32 @@ async def ast_call(self, arg):
             else:
                 kwargs[kw_arg.arg] = await self.aeval(kw_arg.value)
         args = await self.eval_elt_list(arg.args)
-        arg_str = ", ".join(['"' + elt + '"' if isinstance(elt, str) else str(elt) for elt in args])
         #
         # try to deduce function name, although this only works in simple cases
         #
+        func_name = None
         if isinstance(arg.func, ast.Name):
             func_name = arg.func.id
         elif isinstance(arg.func, ast.Attribute):
             func_name = arg.func.attr
-        else:
-            func_name = "<function>"
         if isinstance(func, EvalLocalVar):
             func_name = func.get_name()
             func = func.get()
-        _LOGGER.debug("%s: calling %s(%s, %s)", self.name, func_name, arg_str, kwargs)
         return await self.call_func(func, func_name, *args, **kwargs)
 
     async def call_func(self, func, func_name, *args, **kwargs):
         """Call a function with the given arguments."""
-        if isinstance(func, (EvalFuncVar, EvalFuncVarAstCtx)):
+        if func_name is None:
+            try:
+                if isinstance(func, (EvalFunc, EvalFuncVar, EvalFuncVarAstCtx)):
+                    func_name = func.get_name()
+                else:
+                    func_name = func.__name__
+            except Exception:
+                func_name = "<function>"
+        arg_str = ", ".join(['"' + elt + '"' if isinstance(elt, str) else str(elt) for elt in args])
+        _LOGGER.debug("%s: calling %s(%s, %s)", self.name, func_name, arg_str, kwargs)
+        if isinstance(func, (EvalFunc, EvalFuncVar, EvalFuncVarAstCtx)):
             return await func.call(self, *args, **kwargs)
         if inspect.isclass(func) and hasattr(func, "__init__evalfunc_wrap__"):
             inst = func()

custom_components/pyscript/function.py

@@ -35,6 +35,11 @@ class Function:
     #
     our_tasks = set()
 
+    #
+    # Done callbacks for each task
+    #
+    task2cb = {}
+
     #
     # initial list of available functions
     #
@@ -149,7 +154,6 @@ async def task_waiter(waiter_q):
         if not cls.task_waiter:
             cls.task_waiter_q = asyncio.Queue(0)
             cls.task_waiter = Function.create_task(task_waiter(cls.task_waiter_q))
-            _LOGGER.debug("task_waiter: started")
 
     @classmethod
     def reaper_cancel(cls, task):
@@ -365,7 +369,12 @@ async def run_coro(cls, coro):
         try:
             task = asyncio.current_task()
             cls.our_tasks.add(task)
-            return await coro
+            result = await coro
+            if task in cls.task2cb:
+                for callback, info in cls.task2cb[task]["cb"].items():
+                    ast_ctx, args, kwargs = info
+                    await ast_ctx.call_func(callback, None, *args, **kwargs)
+            return result
         except asyncio.CancelledError:
             raise
         except Exception:
@@ -375,8 +384,8 @@ async def run_coro(cls, coro):
                 for name in cls.unique_task2name[task]:
                     del cls.unique_name2task[name]
                 del cls.unique_task2name[task]
-            if task in cls.task2context:
-                del cls.task2context[task]
+            cls.task2context.pop(task, None)
+            cls.task2cb.pop(task, None)
             cls.our_tasks.discard(task)
 
     @classmethod
@@ -402,3 +411,20 @@ def service_remove(cls, domain, service):
             return
         cls.service_cnt[key] = 0
         cls.hass.services.async_remove(domain, service)
+
+    @classmethod
+    def task_done_callback_ctx(cls, task, ast_ctx):
+        """Set the ast_ctx for a task, which is needed for done callbacks."""
+        cls.task2cb[task] = {"ctx": ast_ctx, "cb": {}}
+
+    @classmethod
+    def task_add_done_callback(cls, task, ast_ctx, callback, *args, **kwargs):
+        """Add a done callback to the given task."""
+        if ast_ctx is None:
+            ast_ctx = cls.task2cb[task]["ctx"]
+        cls.task2cb[task]["cb"][callback] = [ast_ctx, args, kwargs]
+
+    @classmethod
+    def task_remove_done_callback(cls, task, callback):
+        """Remove a done callback to the given task."""
+        cls.task2cb[task]["cb"].pop(callback, None)

custom_components/pyscript/trigger.py

@@ -14,7 +14,7 @@
 import homeassistant.helpers.sun as sun
 
 from .const import LOGGER_PATH
-from .eval import AstEval, EvalFuncVar, EvalFuncVarAstCtx
+from .eval import AstEval, EvalFunc, EvalFuncVar, EvalFuncVarAstCtx
 from .event import Event
 from .function import Function
 from .mqtt import Mqtt
@@ -139,21 +139,28 @@ def user_task_create_factory(ast_ctx):
             async def user_task_create(func, *args, **kwargs):
                 """Implement task.create()."""
 
-                async def func_call(func, new_ast_ctx, *args, **kwargs):
+                async def func_call(func, func_name, new_ast_ctx, *args, **kwargs):
                     """Call user function inside task.create()."""
-                    ret = await func.call(new_ast_ctx, *args, **kwargs)
+                    ret = await new_ast_ctx.call_func(func, func_name, *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()")
+                try:
+                    if isinstance(func, (EvalFunc, EvalFuncVar, EvalFuncVarAstCtx)):
+                        func_name = func.get_name()
+                    else:
+                        func_name = func.__name__
+                except Exception:
+                    func_name = "<function>"
 
                 new_ast_ctx = AstEval(
-                    f"{ast_ctx.get_global_ctx_name()}.{func.get_name()}", ast_ctx.get_global_ctx()
+                    f"{ast_ctx.get_global_ctx_name()}.{func_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))
+                task = Function.create_task(func_call(func, func_name, new_ast_ctx, *args, **kwargs))
+                Function.task_done_callback_ctx(task, new_ast_ctx)
+                return task
 
             return user_task_create
 
@@ -173,9 +180,25 @@ async def user_task_wait(aws):
             """Implement task.wait()."""
             return await asyncio.wait(aws)
 
+        async def user_task_add_done_callback(task, callback, *args, **kwargs):
+            """Implement task.add_done_callback()."""
+            ast_ctx = None
+            if type(callback) is EvalFuncVarAstCtx:
+                ast_ctx = callback.get_ast_ctx()
+                callback = callback.get_eval_func_var()
+            Function.task_add_done_callback(task, ast_ctx, callback, *args, **kwargs)
+
+        async def user_task_remove_done_callback(task, callback):
+            """Implement task.remove_done_callback()."""
+            if type(callback) is EvalFuncVarAstCtx:
+                callback = callback.get_eval_func_var()
+            Function.task_remove_done_callback(task, callback)
+
         funcs = {
             "task.cancel": user_task_cancel,
             "task.wait": user_task_wait,
+            "task.add_done_callback": user_task_add_done_callback,
+            "task.remove_done_callback": user_task_remove_done_callback,
         }
         Function.register(funcs)
 
@@ -1132,7 +1155,7 @@ async def do_func_call(func, ast_ctx, task_unique, task_unique_func, hass_contex
 
             if task_unique and task_unique_func:
                 await task_unique_func(task_unique)
-            await func.call(ast_ctx, **kwargs)
+            await ast_ctx.call_func(func, None, **kwargs)
             if ast_ctx.get_exception_obj():
                 ast_ctx.get_logger().error(ast_ctx.get_exception_long())
 

tests/test_tasks.py

@@ -0,0 +1,149 @@
+"""Test the pyscript apps, modules and import features."""
+
+import asyncio
+import re
+
+from custom_components.pyscript.const import DOMAIN, FOLDER
+from mock_open import MockOpen
+from pytest_homeassistant_custom_component.async_mock import patch
+
+from homeassistant.const import EVENT_STATE_CHANGED
+from homeassistant.setup import async_setup_component
+
+
+async def wait_until_done(notify_q):
+    """Wait for the done handshake."""
+    return await asyncio.wait_for(notify_q.get(), timeout=4)
+
+
+async def test_tasks(hass, caplog):
+    """Test starting tasks."""
+
+    conf_dir = hass.config.path(FOLDER)
+
+    file_contents = {
+        f"{conf_dir}/hello.py": """
+
+#
+# check starting multiple tasks, each stopping the prior one
+#
+def task1(cnt, last):
+    task.unique('task1')
+    if not last:
+        task.sleep(10)
+    log.info(f"finished task1, cnt={cnt}")
+
+for cnt in range(10):
+    task.create(task1, cnt, cnt == 9)
+
+#
+# check the return value after wait
+#
+def task2(arg):
+    return 2 * arg
+
+t2a = task.create(task2, 21)
+t2b = task.create(task2, 51)
+done, pending = task.wait({t2a, t2b})
+log.info(f"task2() results = {[t2a.result(), t2b.result()]}, len(done) = {len(done)};")
+
+#
+# check the return value with a regular function
+#
+@pyscript_compile
+def task3(arg):
+    return 2 * arg
+
+t3a = task.create(task3, 22)
+t3b = task.create(task3, 52)
+done, pending = task.wait({t3a, t3b})
+log.info(f"task3() results = {[t3a.result(), t3b.result()]}, len(done) = {len(done)};")
+
+
+#
+# check that we can do a done callback
+#
+def task4(arg):
+    task.wait_until(state_trigger="pyscript.var4 == '1'")
+    return 2 * arg
+
+def callback4a(arg):
+    log.info(f"callback4a arg = {arg}")
+
+def callback4b(arg):
+    log.info(f"callback4b arg = {arg}")
+
+def callback4c(arg):
+    log.info(f"callback4c arg = {arg}")
+
+t4 = task.create(task4, 23)
+task.add_done_callback(t4, callback4a, 26)
+task.add_done_callback(t4, callback4b, 101)
+task.add_done_callback(t4, callback4c, 200)
+task.add_done_callback(t4, callback4a, 25)
+task.add_done_callback(t4, callback4c, 201)
+task.add_done_callback(t4, callback4b, 100)
+task.add_done_callback(t4, callback4a, 24)
+task.remove_done_callback(t4, callback4c)
+task.remove_done_callback(t4, task4)
+pyscript.var4 = 1
+done, pending = task.wait({t4})
+log.info(f"task4() result = {t4.result()}, len(done) = {len(done)};")
+
+""",
+    }
+
+    mock_open = MockOpen()
+    for key, value in file_contents.items():
+        mock_open[key].read_data = value
+
+    def isfile_side_effect(arg):
+        return arg in file_contents
+
+    def glob_side_effect(path, recursive=None):
+        result = []
+        path_re = path.replace("*", "[^/]*").replace(".", "\\.")
+        path_re = path_re.replace("[^/]*[^/]*/", ".*")
+        for this_path in file_contents:
+            if re.match(path_re, this_path):
+                result.append(this_path)
+        return result
+
+    conf = {"apps": {"world": {}}}
+    with patch("custom_components.pyscript.os.path.isdir", return_value=True), patch(
+        "custom_components.pyscript.glob.iglob"
+    ) as mock_glob, patch("custom_components.pyscript.global_ctx.open", mock_open), patch(
+        "custom_components.pyscript.open", mock_open
+    ), patch(
+        "homeassistant.config.load_yaml_config_file", return_value={"pyscript": conf}
+    ), patch(
+        "custom_components.pyscript.os.path.getmtime", return_value=1000
+    ), patch(
+        "custom_components.pyscript.global_ctx.os.path.getmtime", return_value=1000
+    ), patch(
+        "custom_components.pyscript.os.path.isfile"
+    ) as mock_isfile:
+        mock_isfile.side_effect = isfile_side_effect
+        mock_glob.side_effect = glob_side_effect
+        assert await async_setup_component(hass, "pyscript", {DOMAIN: conf})
+
+    notify_q = asyncio.Queue(0)
+
+    async def state_changed(event):
+        var_name = event.data["entity_id"]
+        if var_name != "pyscript.done":
+            return
+        value = event.data["new_state"].state
+        await notify_q.put(value)
+
+    hass.bus.async_listen(EVENT_STATE_CHANGED, state_changed)
+
+    assert caplog.text.count("finished task1, cnt=9") == 1
+    assert "task2() results = [42, 102], len(done) = 2;" in caplog.text
+    assert "task3() results = [44, 104], len(done) = 2;" in caplog.text
+    assert "task4() result = 46, len(done) = 1;" in caplog.text
+    assert caplog.text.count("callback4a arg =") == 1
+    assert "callback4a arg = 24" in caplog.text
+    assert caplog.text.count("callback4b arg =") == 1
+    assert "callback4b arg = 100" in caplog.text
+    assert "callback4c arg =" not in caplog.text