From f8132a7978722d9b0287842761bb3dc450c36e59 Mon Sep 17 00:00:00 2001 From: gnzsnz <8376642+gnzsnz@users.noreply.github.com> Date: Thu, 25 Sep 2025 18:17:51 +0200 Subject: [PATCH 1/3] Enhance event loop handling in Event class and utility functions - address #8 - Added check to prevent running Event.run() within an already running asyncio event loop. - Implemented cleanup of pending tasks after running the event list. - Updated get_event_loop() to handle different scenarios for retrieving the asyncio event loop. --- eventkit/event.py | 31 ++++++++++++++++++++++++++++++- eventkit/util.py | 13 +++++++++++-- 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/eventkit/event.py b/eventkit/event.py index a3b405b..22ce345 100644 --- a/eventkit/event.py +++ b/eventkit/event.py @@ -346,7 +346,36 @@ 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. + result = await self.list() + + # 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"): """ diff --git a/eventkit/util.py b/eventkit/util.py index da34ab9..c704a28 100644 --- a/eventkit/util.py +++ b/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]: From 0ae4dcdfac110639efd11425c9ed5951eef49659 Mon Sep 17 00:00:00 2001 From: gnzsnz <8376642+gnzsnz@users.noreply.github.com> Date: Sat, 27 Sep 2025 17:21:14 +0200 Subject: [PATCH 2/3] Implement cancellation handling in event operations and add comprehensive cancellation tests --- eventkit/event.py | 29 +++- eventkit/ops/create.py | 46 +++-- eventkit/ops/op.py | 7 + eventkit/ops/transform.py | 111 +++++++++---- pyproject.toml | 4 +- tests/cancellation_test.py | 333 +++++++++++++++++++++++++++++++++++++ 6 files changed, 483 insertions(+), 47 deletions(-) create mode 100644 tests/cancellation_test.py diff --git a/eventkit/event.py b/eventkit/event.py index 22ce345..5d874d2 100644 --- a/eventkit/event.py +++ b/eventkit/event.py @@ -355,7 +355,14 @@ def run(self) -> List: async def _run_and_clean(): # This coroutine will be run to completion. - result = await self.list() + 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() @@ -519,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 @@ -593,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. diff --git a/eventkit/ops/create.py b/eventkit/ops/create.py index 24774b3..dbd9ad3 100644 --- a/eventkit/ops/create.py +++ b/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() diff --git a/eventkit/ops/op.py b/eventkit/ops/op.py index 1927d33..71bdd9d 100644 --- a/eventkit/ops/op.py +++ b/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: diff --git a/eventkit/ops/transform.py b/eventkit/ops/transform.py index afcdd38..bc742a0 100644 --- a/eventkit/ops/transform.py +++ b/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): diff --git a/pyproject.toml b/pyproject.toml index 4d9f58f..9b8c21f 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/cancellation_test.py b/tests/cancellation_test.py new file mode 100644 index 0000000..f963e3e --- /dev/null +++ b/tests/cancellation_test.py @@ -0,0 +1,333 @@ +"""Cancellation Tests""" + +import asyncio +import logging + +import pytest + +import eventkit as ev + +logger = logging.getLogger(__name__) + + +@pytest.mark.asyncio +async def test_aiterate_cancellation(): + """Verify that breaking an async for loop cancels the source Aiterate task.""" + generator_finished_cleanly = False + + async def my_generator(): + nonlocal generator_finished_cleanly + try: + for i in range(10): + yield i + await asyncio.sleep(0.1) + finally: + generator_finished_cleanly = True + + event = ev.Event.aiterate(my_generator()) + count = 0 + async for _ in event: + count += 1 + if count == 2: + break + await asyncio.sleep(0.1) + + assert generator_finished_cleanly, "Generator was not cancelled correctly." + + +@pytest.mark.asyncio +async def test_wait_cancellation(): + """Verify that cancelling a Wait event cancels the underlying Future.""" + long_running_future = asyncio.Future() + wait_event = ev.Event.wait(long_running_future) + assert not long_running_future.done() + wait_event.cancel() + await asyncio.sleep(0) + assert long_running_future.cancelled() + + +@pytest.mark.asyncio +async def test_timer_cancellation(): + """Verify that breaking an async for loop cancels the source Timer task.""" + # Use ev.Timer directly + interval = 0.1 + count_limit = 10 # The timer will try to emit 10 times + num_to_take = 2 # We will take 2 items, triggering cancellation + + timer_event = ev.Timer(interval, count=count_limit) + + results = [] + try: + async for item in timer_event.take(num_to_take): + results.append(item) + await asyncio.sleep(0) # Yield control + except asyncio.CancelledError: + pass + + assert len(results) == num_to_take + + # Give some time for cancellation to propagate + await asyncio.sleep(interval * (count_limit - num_to_take) + 0.1) + + # The timer should be done after cancellation + assert timer_event.done() + + +@pytest.mark.asyncio +async def test_map_cancellation(): + """Verify that Map cancels its in-flight tasks.""" + processed_items = [] + tasks_cancelled = [] + + async def long_running_map(x): + try: + processed_items.append(x) + # Make the first two items fast and the rest slow + sleep_time = 0.1 if x < 2 else 1 + await asyncio.sleep(sleep_time) + return x + except asyncio.CancelledError: + tasks_cancelled.append(x) + raise + + event = ev.Range(10).map(long_running_map, task_limit=5).take(2) + results = [] + try: + async for item in event: + results.append(item) + except asyncio.CancelledError: + pass + + assert len(results) == 2 + # Wait long enough for the slow tasks to be cancelled + await asyncio.sleep(1.5) + + assert len(processed_items) < 6 + assert len(processed_items) > 1 + assert len(tasks_cancelled) == 3 + + +@pytest.mark.asyncio +async def test_map_cancellation_unordered(): + """Verify that Map(ordered=False) cancels its in-flight tasks.""" + processed_items = [] + tasks_cancelled = [] + + async def long_running_map(x): + try: + processed_items.append(x) + # Make the first two items fast and the rest slow + sleep_time = 0.1 if x < 2 else 1.0 + await asyncio.sleep(sleep_time) + return x + except asyncio.CancelledError: + tasks_cancelled.append(x) + raise + + event = ev.Range(10).map(long_running_map, task_limit=5, ordered=False).take(2) + results = [] + try: + async for item in event: + results.append(item) + await asyncio.sleep(0) # Yield control to event loop + except asyncio.CancelledError: + pass + + assert len(results) == 2 + await asyncio.sleep(1.5) + + assert len(processed_items) <= 6 + assert len(processed_items) > 1 + assert len(tasks_cancelled) == 3 + + +@pytest.mark.asyncio +async def test_map_cancellation_general_case(): + """Verify general cancellation behavior for unordered Map with multiple tasks.""" + processed_items = [] + tasks_cancelled = [] + num_to_take = 3 + total_items = 10 + task_limit = 5 + + async def long_running_map(x): + try: + processed_items.append(x) + # Make the first num_to_take items fast and the rest slow + sleep_time = 0.1 if x < num_to_take else 1.0 + await asyncio.sleep(sleep_time) + return x + except asyncio.CancelledError: + tasks_cancelled.append(x) + raise + + event = ( + ev.Range(total_items) + .map(long_running_map, task_limit=task_limit, ordered=False) + .take(num_to_take) + ) + results = [] + try: + async for item in event: + results.append(item) + await asyncio.sleep(0) # Yield control to event loop + except asyncio.CancelledError: + pass + + assert len(results) == num_to_take + + # Wait for cancellation to propagate + await asyncio.sleep(1.5) + + # Expected behavior: + # - num_to_take items are received. + # - task_limit tasks are started initially. + # - Some tasks might complete before cancellation. + # - The remaining in-flight tasks should be cancelled. + # - The number of processed_items should be num_to_take + (tasks that started + # before cancellation but were not cancelled) + # - The number of tasks_cancelled should be the remaining in-flight tasks. + + # A more robust assertion for processed_items: + # It should be at least num_to_take (the ones we received) + # And at most task_limit + num_to_take (the ones that could have started before + # cancellation) + assert len(processed_items) >= num_to_take + assert ( + len(processed_items) <= task_limit + num_to_take + ) # Max possible started tasks + + # We expect a significant number of tasks to be cancelled. + # The number of tasks that were started but not completed and then cancelled. + # This should be roughly total_items - num_to_take - (tasks that completed before + # cancellation) + assert len(tasks_cancelled) > 0 + assert ( + len(tasks_cancelled) <= total_items - num_to_take + ) # Max possible cancelled tasks + + +@pytest.mark.asyncio +async def test_chained_map_cancellation(): + """Verify cancellation propagates through chained Map operators.""" + processed_items_1 = [] + tasks_cancelled_1 = [] + processed_items_2 = [] + tasks_cancelled_2 = [] + + num_to_take = 2 + total_items = 10 + task_limit_1 = 5 + task_limit_2 = 3 + + async def long_running_map_1(x): + try: + processed_items_1.append(x) + # Make the first num_to_take items fast, the rest slow + sleep_time = 0.1 if x < num_to_take else 1.0 + await asyncio.sleep(sleep_time) + return x + except asyncio.CancelledError: + tasks_cancelled_1.append(x) + raise + + async def long_running_map_2(x): + try: + processed_items_2.append(x) + # Make all tasks in Map2 slow + sleep_time = 1.0 + await asyncio.sleep(sleep_time) + return x + except asyncio.CancelledError: + tasks_cancelled_2.append(x) + raise + + event = ( + ev.Range(total_items) + .map(long_running_map_1, task_limit=task_limit_1, ordered=False) + .map(long_running_map_2, task_limit=task_limit_2, ordered=False) + .take(num_to_take) + ) + results = [] + try: + async for item in event: + results.append(item) + await asyncio.sleep(0) # Yield control to event loop + except asyncio.CancelledError: + pass + + assert len(results) == num_to_take + + # Wait for cancellation to propagate through both layers + await asyncio.sleep(1.5) + + # Assertions for Map1 + assert len(processed_items_1) >= num_to_take + assert len(processed_items_1) <= total_items + assert len(tasks_cancelled_1) > 0 + assert len(tasks_cancelled_1) <= total_items - num_to_take + + # Assertions for Map2 + assert len(processed_items_2) >= num_to_take + assert len(processed_items_2) <= total_items + assert len(tasks_cancelled_2) > 0 + assert len(tasks_cancelled_2) <= total_items - num_to_take + + +@pytest.mark.asyncio +async def test_aiterate_source_map_cancellation(): + """Verify cancellation propagates from Map back to an Aiterate source.""" + generator_finished_cleanly = False + processed_items = [] + tasks_cancelled = [] + + num_to_take = 3 + total_items = 10 + task_limit = 5 + + async def my_generator(): + nonlocal generator_finished_cleanly + try: + for i in range(total_items): + yield i + await asyncio.sleep(0.1) # Simulate some work + finally: + generator_finished_cleanly = True + + async def long_running_map(x): + try: + processed_items.append(x) + # Make the first num_to_take items fast, the rest slow + sleep_time = 0.1 if x < num_to_take else 1.0 + await asyncio.sleep(sleep_time) + return x + except asyncio.CancelledError: + tasks_cancelled.append(x) + raise + + event = ( + ev.Event.aiterate(my_generator()) + .map(long_running_map, task_limit=task_limit, ordered=False) + .take(num_to_take) + ) + results = [] + try: + async for item in event: + results.append(item) + await asyncio.sleep(0) # Yield control to event loop + except asyncio.CancelledError: + pass + + assert len(results) == num_to_take + + # Wait for cancellation to propagate + await asyncio.sleep(1.5) + + # Assertions for Map + assert len(processed_items) >= num_to_take + assert len(processed_items) <= total_items + assert len(tasks_cancelled) > 0 + assert len(tasks_cancelled) <= total_items - num_to_take + + # Assertions for Aiterate source + assert generator_finished_cleanly, "Aiterate generator was not cancelled cleanly." From 178550bc588586cc31ec7c49bdd8b33fe3a2cf41 Mon Sep 17 00:00:00 2001 From: gnzsnz <8376642+gnzsnz@users.noreply.github.com> Date: Sat, 25 Oct 2025 19:02:38 +0200 Subject: [PATCH 3/3] fix for #12 Fix Timeout class to call parent on_source method and enhance timeout tests --- eventkit/ops/timing.py | 1 + tests/timing_test.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/eventkit/ops/timing.py b/eventkit/ops/timing.py index 844368f..a9fb770 100644 --- a/eventkit/ops/timing.py +++ b/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() diff --git a/tests/timing_test.py b/tests/timing_test.py index 77cbbe3..c47391a 100644 --- a/tests/timing_test.py +++ b/tests/timing_test.py @@ -23,6 +23,13 @@ def test_sample(self): self.assertEqual(event.run(), [2, 4, 6, 8]) def test_timeout(self): + # source faster than timeout + seq = Event.sequence(array1, interval=0.01).timeout(0.1) + self.assertEqual(seq.run(), array1) + # source slower than timeout + seq2 = Event.sequence([1, 2, 3], interval=0.1).timeout(0.01) + self.assertEqual(seq2.run(), [1, Event.NO_VALUE]) + # plain timeout timer = Event.timer(10, count=1) event = timer.timeout(0.01) self.assertEqual(event.run(), [Event.NO_VALUE])