task.unique() doesn't kill self if self previously called it

Author: craigbarratt

custom_components/pyscript/function.py

@@ -107,28 +107,28 @@ def task_unique_factory(cls, ctx):
         async def task_unique(name, kill_me=False):
             """Implement task.unique()."""
             name = f"{ctx.get_global_ctx_name()}.{name}"
+            curr_task = asyncio.current_task()
             if name in cls.unique_name2task:
+                task = cls.unique_name2task[name]
                 if kill_me:
-                    #
-                    # it seems we can't cancel ourselves, so we
-                    # tell the repeaer task to cancel us
-                    #
-                    Function.task_cancel(asyncio.current_task())
-                    # wait to be canceled
-                    await asyncio.sleep(100000)
-                else:
-                    task = cls.unique_name2task[name]
-                    if task in cls.our_tasks:
-                        # only cancel tasks if they are ones we started
-                        try:
-                            task.cancel()
-                            await task
-                        except asyncio.CancelledError:
-                            pass
-            task = asyncio.current_task()
-            if task in cls.our_tasks:
-                cls.unique_name2task[name] = task
-                cls.unique_task2name[task] = name
+                    if task != curr_task:
+                        #
+                        # it seems we can't cancel ourselves, so we
+                        # tell the repeaer task to cancel us
+                        #
+                        Function.task_cancel(curr_task)
+                        # wait to be canceled
+                        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
+            if curr_task in cls.our_tasks:
+                cls.unique_name2task[name] = curr_task
+                cls.unique_task2name[curr_task] = name
 
         return task_unique
 

tests/test_unique.py

@@ -75,6 +75,11 @@ def funcStartupSync():
     seq_num += 1
     log.info(f"funcStartupSync setting pyscript.done = {seq_num}")
     pyscript.done = seq_num
+    #
+    # stick around so the task.unique() still applies
+    #
+    task.unique("func6")
+    task.sleep(10000)
 
 @state_trigger("pyscript.f1var1 == '1'")
 def func1(var_name=None, value=None):
@@ -132,9 +137,8 @@ def func5():
     pyscript.done = [seq_num, "pyscript.f5var2"]
 
 
-@time_trigger("startup")
+@state_trigger("pyscript.f4var1 == '1'")
 def func6():
-    task.unique("func6")
     task.unique("func6", kill_me=True)
     # mess up the sequence numbers if task.unique fails to kill us
     pyscript.done = [999]