add timeout option in tasks.create

nggit · Nov 7, 2023 · 852a613 · 852a613
1 parent 6ec381e
commit 852a613
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/tremolo/lib/tasks.py b/tremolo/lib/tasks.py
@@ -11,8 +11,13 @@ def __init__(self, tasks, loop=None):
         self._loop = loop
         self._tasks = tasks
 
-    def create(self, coro):
+    def create(self, coro, timeout=0):
         task = self._loop.create_task(coro)
-        self._tasks.append(task.cancel)
+
+        if timeout > 0:
+            self._loop.call_at(self._loop.time() + timeout, task.cancel)
+        else:
+            # until the connection is lost
+            self._tasks.append(task.cancel)
 
         return task