Bug timeout

eventkit/event.py

@@ -346,7 +346,43 @@ def run(self) -> List:
                 await event.list()
         """
         loop = get_event_loop()
-        return loop.run_until_complete(self.list())
+
+        if loop.is_running():
+            raise RuntimeError(
+                "Event.run() cannot be called from within an already running "
+                "asyncio event loop. Use 'await event.list()' instead."
+            )
+
+        async def _run_and_clean():
+            # This coroutine will be run to completion.
+            list_op_event = self.list()  # Get the ListOp event
+            result = (
+                await list_op_event
+            )  # Await its completion (which means it emits its list)
+
+            # Explicitly wait for the ListOp event to be done.
+            # This ensures its source (e.g., Aiterate) has also completed its lifecycle.
+            await list_op_event.wait_until_done()
+
+            # Now, perform the explicit cleanup of any other background tasks.
+            current_loop = asyncio.get_running_loop()
+            all_tasks = asyncio.all_tasks(loop=current_loop)
+            current_task = asyncio.current_task(loop=current_loop)
+
+            pending_tasks = [
+                task
+                for task in all_tasks
+                if task is not current_task and not task.done()
+            ]
+
+            if pending_tasks:
+                for task in pending_tasks:
+                    task.cancel()
+                await asyncio.gather(*pending_tasks, return_exceptions=True)
+
+            return result
+
+        return loop.run_until_complete(_run_and_clean())
 
     def pipe(self, *targets: "Event"):
         """
@@ -490,6 +526,8 @@ def on_done(source):
                     break
         finally:
             self.disconnect(on_event, on_error, on_done)
+            if hasattr(self, "cancel"):
+                self.cancel()
 
     __iadd__ = connect
     __isub__ = disconnect
@@ -564,6 +602,24 @@ def __contains__(self, c):
         obj, func = self._split(c)
         return self._slots.exists(obj, func)
 
+    async def wait_until_done(self):
+        """
+        Asynchronously waits until this event is done.
+        """
+        if self.done():
+            return
+        fut = asyncio.Future()
+
+        def on_done(source):
+            if not fut.done():
+                fut.set_result(None)
+
+        self.done_event.connect(on_done)
+        try:
+            await fut
+        finally:
+            self.done_event.disconnect(on_done)
+
     def __reduce__(self):
         """
         Don't pickle slots.

eventkit/ops/create.py

@@ -1,11 +1,14 @@
 import asyncio
 import itertools
+import logging
 import time
 
 from ..event import Event
 from ..util import NO_VALUE, get_event_loop, timerange
 from .op import Op
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 class Wait(Event):
     __slots__ = ("_task",)
@@ -16,8 +19,8 @@ def __init__(self, future, name="wait"):
             self._task = None
             self.set_done()
         else:
-            # Note: the loop= *is* necessary here.
-            self._task = asyncio.ensure_future(future, loop=get_event_loop())
+            loop = get_event_loop()
+            self._task = asyncio.ensure_future(future, loop=loop)
             future.add_done_callback(self._on_task_done)
 
     def _on_task_done(self, task):
@@ -35,27 +38,50 @@ def __del__(self):
         if self._task:
             self._task.cancel()
 
+    def cancel(self):
+        if self._task and not self._task.done():
+            self._task.cancel()
+
 
 class Aiterate(Event):
-    __slots__ = ("_task",)
+    __slots__ = ("_task", "_ait")
 
     def __init__(self, ait):
         Event.__init__(self, ait.__qualname__)
+        self._ait = ait
+        self._task = None
 
-        # Note: the loop= *is* necessary here.
-        self._task = asyncio.ensure_future(self._looper(ait), loop=get_event_loop())
-
-    async def _looper(self, ait):
+    async def _looper(self):
         try:
-            async for args in ait:
+            async for args in self._ait:
                 self.emit(args)
+        except asyncio.CancelledError:
+            # Task was cancelled, clean up and exit gracefully
+            logger.debug(f"Aiterate task for {self.name()} cancelled.")
         except Exception as error:
             self.error_event.emit(self, error)
+        finally:
+            self.set_done()
 
-        self._task = None
-        self.set_done()
+    def connect(self, *args, **kwargs):
+        if self._task is None:
+            loop = get_event_loop()
+            self._task = loop.create_task(self._looper())
+        return super().connect(*args, **kwargs)
+
+    def cancel(self):
+        """
+        Explicitly cancels the underlying asyncio task and sets the event as done.
+        """
+        if self._task and not self._task.done():
+            self._task.cancel()
+            logger.debug(f"Aiterate task for {self.name()} requested cancellation.")
+        self.set_done()  # Re-enabled to fix race condition
 
     def __del__(self):
+        # The __del__ method is unreliable for timely resource cleanup.
+        # Explicit `cancel()` method should be preferred.
+        # However, keeping this for robustness in case `cancel()` is not called.
         if self._task:
             self._task.cancel()
 

eventkit/ops/op.py

@@ -38,6 +38,13 @@ def on_source_done(self, _source):
             self._source = None
         self.set_done()
 
+    def set_done(self):
+        if not self.done():
+            super().set_done()
+            # An operator being done means it no longer needs its source.
+            if self._source and hasattr(self._source, "cancel"):
+                self._source.cancel()
+
     def set_source(self, source):
         source = Event.create(source)
         if self._source is None:

eventkit/ops/timing.py

@@ -44,6 +44,7 @@ def __init__(self, timeout, source=None):
     def on_source(self, *args):
         loop = get_event_loop()
         self._last_time = loop.time()
+        Op.on_source(self, *args)
 
     def on_source_done(self, source):
         self._handle.cancel()

eventkit/ops/transform.py

@@ -1,12 +1,15 @@
 import asyncio
 import copy
+import logging
 import time
 from collections import deque
 
-from ..util import NO_VALUE
+from ..util import NO_VALUE, get_event_loop
 from .combine import Chain, Concat, Merge, Switch
 from .op import Op
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 class Constant(Op):
     __slots__ = ("_constant",)
@@ -210,7 +213,15 @@ def on_source_done(self, source):
 
 
 class Map(Op):
-    __slots__ = ("_func", "_timeout", "_ordered", "_task_limit", "_coro_q", "_tasks")
+    __slots__ = (
+        "_func",
+        "_timeout",
+        "_ordered",
+        "_task_limit",
+        "_coro_q",
+        "_tasks",
+        "_pending_ordered",
+    )
 
     def __init__(self, func, timeout=0, ordered=True, task_limit=None, source=None):
         Op.__init__(self, source)
@@ -222,9 +233,12 @@ def __init__(self, func, timeout=0, ordered=True, task_limit=None, source=None):
         self._ordered = ordered
         self._task_limit = task_limit
         self._coro_q = deque()
-        self._tasks = deque()
+        self._tasks = set()  # the use of `set` eliminates race conditions
+        self._pending_ordered = deque() if ordered else None
 
     def on_source(self, *args):
+        if self.done():
+            return
         obj = self._func(*args)
         if asyncio.iscoroutine(obj):
             # function returns an awaitable
@@ -239,50 +253,77 @@ def on_source(self, *args):
             self.emit(obj)
 
     def on_source_done(self, source):
-        if not self._tasks:
-            # only end when no tasks are pending
-            Op.on_source_done(self, self._source)
-
-        self._source = None
+        if not self._tasks and not self._coro_q:
+            super().on_source_done(source)
 
     def _create_task(self, coro):
-        # schedule a task to be run
         if self._timeout:
             coro = asyncio.wait_for(coro, self._timeout)
 
-        task = asyncio.create_task(coro)
+        loop = get_event_loop()
+        task = loop.create_task(coro)
+        self._tasks.add(task)
         task.add_done_callback(self._on_task_done)
-        self._tasks.append(task)
-
-    def _on_task_done(self, task):
-        # handle task result
-        tasks = self._tasks
         if self._ordered:
-            while tasks and tasks[0].done():
-                # remove task after emitting result
-                task = tasks[0]
-                self._emit_task(task)
-                task = tasks.popleft()
-        else:
-            # remove task after emitting result
-            self._emit_task(task)
-            tasks.remove(task)
-
-        # schedule pending awaitables from the queue
-        while self._coro_q and (not self._task_limit or len(tasks) < self._task_limit):
-            self._create_task(self._coro_q.popleft())
-
-        # end when source has ended with no pending tasks
-        if not tasks and self._source is None:
-            Op.on_source_done(self, self._source)
+            self._pending_ordered.append(task)
 
-    def _emit_task(self, task):
+    def _on_task_done(self, task):
+        """handle task completion"""
+        self._tasks.discard(task)
+
+        if not self.done():  # If we are not in a cancellation flow
+            if self._ordered:
+                while self._pending_ordered and self._pending_ordered[0].done():
+                    done_task = self._pending_ordered.popleft()
+                    self._emit_task_result(done_task)
+            else:
+                self._emit_task_result(task)
+
+            # Schedule new tasks from the queue if space is available
+            while self._coro_q and (
+                not self._task_limit or len(self._tasks) < self._task_limit
+            ):
+                self._create_task(self._coro_q.popleft())
+
+        # Check if everything is finished
+        if (
+            self._source
+            and self._source.done()
+            and not self._tasks
+            and not self._coro_q
+        ):
+            super().on_source_done(self._source)
+
+    def _emit_task_result(self, task):
         try:
             result = task.result()
+            self.emit(result)
+        except asyncio.CancelledError:
+            return
         except Exception as error:
-            result = NO_VALUE
             self.error_event.emit(error)
-        self.emit(result)
+
+    def cancel(self):
+        if self.done():
+            return
+
+        # 1. Immediately mark as done to prevent any more processing
+        super().set_done()
+
+        # 2. Disconnect from source
+        if self._source:
+            self._disconnect_from(self._source)
+            self._source = None
+
+        # 3. Cancel pending work
+        for coro in self._coro_q:
+            coro.close()  # Prevent ResourceWarning: coroutine was never awaited
+        self._coro_q.clear()
+
+        # 4. Cancel all in-flight tasks
+        for task in self._tasks:
+            task.cancel()
+        self._tasks.clear()
 
 
 class Emap(Op):

eventkit/util.py

@@ -1,3 +1,5 @@
+"""Eventkit utilities."""
+
 import asyncio
 import datetime as dt
 from typing import AsyncIterator, Final
@@ -18,8 +20,15 @@ def __repr__(self):
 
 def get_event_loop():
     """Get asyncio event loop or create one if it doesn't exist."""
-    loop = asyncio.get_event_loop_policy().get_event_loop()
-    return loop
+    try:
+        return asyncio.get_running_loop()
+    except RuntimeError:
+        try:
+            return asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop
 
 
 async def timerange(start=0, end=None, step: float = 1) -> AsyncIterator[dt.datetime]:

pyproject.toml

@@ -72,4 +72,6 @@ ignore = [
 ]
 [tool.pytest.ini_options]
 log_cli = true
-asyncio_default_test_loop_scope = "session"
+asyncio_mode = "strict"
+asyncio_default_test_loop_scope = "function"
+asyncio_default_fixture_loop_scope = "function"