diff --git a/extmod/asyncio/core.py b/extmod/asyncio/core.py
index 8aad234514b0c..e53383e7ffac0 100644
--- a/extmod/asyncio/core.py
+++ b/extmod/asyncio/core.py
@@ -129,6 +129,15 @@ def wait_io_event(self, dt):
else:
self.poller.modify(s, select.POLLIN)
+ def all_tasks(self):
+ tasks = set()
+ for q1, q2, _ in self.map.values():
+ if q1 is not None:
+ tasks.add(q1)
+ if q2 is not None:
+ tasks.add(q2)
+
+ return tasks
################################################################################
# Main run loop
@@ -233,7 +242,11 @@ def run_until_complete(main_task=None):
# Create a new task from a coroutine and run it until it finishes
def run(coro):
- return run_until_complete(create_task(coro))
+ try:
+ return run_until_complete(create_task(coro))
+ finally:
+ Loop.close()
+ run_until_complete()
################################################################################
@@ -271,7 +284,17 @@ def stop():
_stop_task = None
def close():
- pass
+ for t in Loop.all_tasks():
+ if not t.done() and t is not cur_task:
+ t.cancel()
+
+ def all_tasks():
+ tasks = set(_task_queue)
+ tasks.update(_io_queue.all_tasks())
+ if cur_task is not None:
+ tasks.add(cur_task)
+
+ return tasks
def set_exception_handler(handler):
Loop._exc_handler = handler
@@ -299,6 +322,13 @@ def current_task():
return cur_task
+def all_tasks():
+ if cur_task is None:
+ raise RuntimeError("no running event loop")
+
+ return Loop.all_tasks()
+
+
def new_event_loop():
global _task_queue, _io_queue
# TaskQueue of Task instances
diff --git a/extmod/asyncio/task.py b/extmod/asyncio/task.py
index 30be2170eb0a7..a9341aa628e66 100644
--- a/extmod/asyncio/task.py
+++ b/extmod/asyncio/task.py
@@ -116,6 +116,12 @@ def pop(self):
def remove(self, v):
self.heap = ph_delete(self.heap, v)
+ def __iter__(self):
+ heap = self.heap
+ while heap:
+ yield heap
+ heap = heap.ph_child
+
# Task class representing a coroutine, can be waited on and cancelled.
class Task:
diff --git a/extmod/modasyncio.c b/extmod/modasyncio.c
index 6161500bf5e7c..c474570037010 100644
--- a/extmod/modasyncio.c
+++ b/extmod/modasyncio.c
@@ -151,6 +151,38 @@ static mp_obj_t task_queue_remove(mp_obj_t self_in, mp_obj_t task_in) {
}
static MP_DEFINE_CONST_FUN_OBJ_2(task_queue_remove_obj, task_queue_remove);
+typedef struct _mp_obj_task_queue_it_t {
+ mp_obj_base_t base;
+ mp_fun_1_t iternext;
+ mp_obj_task_t *cur;
+} mp_obj_task_queue_it_t;
+
+static mp_obj_t task_queue_iternext(mp_obj_t self_in) {
+ mp_obj_task_queue_it_t *self = MP_OBJ_TO_PTR(self_in);
+
+ if (self->cur != NULL) {
+ mp_obj_task_t *current = self->cur;
+ self->cur = (mp_obj_task_t *)self->cur->pairheap.child;
+ return current;
+ }
+ return MP_OBJ_STOP_ITERATION;
+}
+
+static mp_obj_t task_queue_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
+ assert(sizeof(mp_obj_task_queue_it_t) <= sizeof(mp_obj_iter_buf_t));
+ mp_obj_task_queue_t *self = MP_OBJ_TO_PTR(self_in);
+ mp_obj_task_queue_it_t *o = (mp_obj_task_queue_it_t *)iter_buf;
+ o->base.type = &mp_type_polymorph_iter;
+ o->iternext = task_queue_iternext;
+ o->cur = self->heap;
+ return MP_OBJ_FROM_PTR(o);
+}
+
+static const mp_getiter_iternext_custom_t task_queue_getiter_iternext = {
+ .getiter = task_queue_getiter,
+ .iternext = task_queue_iternext,
+};
+
static const mp_rom_map_elem_t task_queue_locals_dict_table[] = {
{ MP_ROM_QSTR(MP_QSTR_peek), MP_ROM_PTR(&task_queue_peek_obj) },
{ MP_ROM_QSTR(MP_QSTR_push), MP_ROM_PTR(&task_queue_push_obj) },
@@ -162,9 +194,10 @@ static MP_DEFINE_CONST_DICT(task_queue_locals_dict, task_queue_locals_dict_table
static MP_DEFINE_CONST_OBJ_TYPE(
task_queue_type,
MP_QSTR_TaskQueue,
- MP_TYPE_FLAG_NONE,
+ MP_TYPE_FLAG_ITER_IS_CUSTOM,
make_new, task_queue_make_new,
- locals_dict, &task_queue_locals_dict
+ locals_dict, &task_queue_locals_dict,
+ iter, &task_queue_getiter_iternext
);
/******************************************************************************/
diff --git a/tests/extmod/asyncio_all_tasks.py b/tests/extmod/asyncio_all_tasks.py
new file mode 100644
index 0000000000000..54a00b49e1664
--- /dev/null
+++ b/tests/extmod/asyncio_all_tasks.py
@@ -0,0 +1,50 @@
+try:
+ import asyncio
+except ImportError:
+ print("SKIP")
+ raise SystemExit
+
+import sys
+
+async def waiter():
+ await asyncio.sleep(1)
+
+async def reader():
+ await asyncio.StreamReader().read()
+
+
+async def print_all_tasks():
+ print(list('T' for t in asyncio.all_tasks()))
+
+async def short_wait():
+ await asyncio.sleep(0.01)
+
+async def waiters():
+ t1 = asyncio.create_task(waiter())
+ # NOTE: The singleton generator for ::sleep expects the task to be scheduled
+ # before using it again.
+ await asyncio.sleep(0)
+ t2 = asyncio.create_task(waiter())
+ await asyncio.sleep(0)
+ t3 = asyncio.create_task(short_wait())
+ # Tasks are in the list before they start to run
+ assert t1 in asyncio.all_tasks()
+ assert t2 in asyncio.all_tasks()
+ await print_all_tasks()
+ await t3
+ assert t1 in asyncio.all_tasks()
+ assert t2 in asyncio.all_tasks()
+ assert t3 not in asyncio.all_tasks()
+
+async def readers():
+ t1 = asyncio.create_task(reader())
+ t2 = asyncio.create_task(reader())
+ await print_all_tasks()
+ assert t1 in asyncio.all_tasks()
+ assert t2 in asyncio.all_tasks()
+
+asyncio.run(print_all_tasks())
+asyncio.new_event_loop()
+asyncio.run(readers())
+asyncio.new_event_loop()
+asyncio.run(waiters())
\ No newline at end of file
diff --git a/tests/extmod/asyncio_shutdown.py b/tests/extmod/asyncio_shutdown.py
new file mode 100644
index 0000000000000..07285977c113a
--- /dev/null
+++ b/tests/extmod/asyncio_shutdown.py
@@ -0,0 +1,135 @@
+try:
+ import asyncio
+except ImportError:
+ print("SKIP")
+ raise SystemExit
+
+import sys
+
+# CPython and Micropython disagree on the sorting of tasks in the ::all_tasks
+# set, so sort the output of the tasks as they shut down rather than expecting
+# them to shut down in the same order.
+class SortedList(list):
+ def print(self):
+ print(*sorted(self))
+
+finished = SortedList()
+
+async def task1():
+ try:
+ await asyncio.sleep(1)
+ finally:
+ finished.append('task1 finished')
+
+async def task2():
+ try:
+ await asyncio.sleep(1)
+ finally:
+ finished.append("task2 finished")
+
+async def reader():
+ try:
+ await asyncio.StreamReader(sys.stdin).read()
+ finally:
+ finished.append("reader finished")
+
+# KeyboardInterrupt in main task should be propagated
+async def ctrlc():
+ asyncio.create_task(task1())
+ await asyncio.sleep(0)
+ asyncio.create_task(task2())
+ await asyncio.sleep(0.01)
+ raise KeyboardInterrupt
+
+
+# If the main task is canceled, the loop should be closed
+async def cancel_main():
+ async def canceler(delay, task):
+ await asyncio.sleep(delay)
+ task.cancel()
+
+ asyncio.create_task(task1())
+ await asyncio.sleep(0)
+ asyncio.create_task(task2())
+ asyncio.create_task(canceler(0.01, asyncio.current_task()))
+ await asyncio.sleep(0.02)
+
+# SystemExit should close the loop
+async def sys_exit():
+ asyncio.create_task(task1())
+ await asyncio.sleep(0)
+ asyncio.create_task(task2())
+ await asyncio.sleep(0.01)
+ raise SystemExit
+
+# Cancelling a background task should have no effect
+async def cancel_bg():
+ t1 = asyncio.create_task(task1())
+ await asyncio.sleep(0)
+ t2 = asyncio.create_task(task2())
+ await asyncio.sleep(0.01)
+ t1.cancel()
+ t2.cancel()
+ await asyncio.sleep(0.01)
+
+# Reader tasks should also be cancelled when the loop is shutdown
+async def cancel_reader():
+ t1 = asyncio.create_task(task1())
+ await asyncio.sleep(0)
+ t2 = asyncio.create_task(reader())
+ await asyncio.sleep(0.01)
+ t1.cancel()
+ t2.cancel()
+ await asyncio.sleep(0.01)
+
+async def close_loop():
+ asyncio.create_task(task1())
+ await asyncio.sleep(0)
+ asyncio.create_task(reader())
+ await asyncio.sleep(0.01)
+
+try:
+ asyncio.new_event_loop()
+ finished.clear()
+ asyncio.run(ctrlc())
+except KeyboardInterrupt:
+ print("KeyboardInterrupt raised")
+finally:
+ finished.print()
+
+try:
+ asyncio.new_event_loop()
+ finished.clear()
+ asyncio.run(cancel_main())
+except asyncio.CancelledError:
+ print("CancelledError raised")
+finally:
+ finished.print()
+
+# Cancel a background waiting task doesn't raise an exception
+print('---')
+asyncio.new_event_loop()
+finished.clear()
+asyncio.run(cancel_bg())
+finished.print()
+
+# Cancel a background IO task doesn't raise an exception
+print('---')
+asyncio.new_event_loop()
+finished.clear()
+asyncio.run(cancel_reader())
+finished.print()
+
+print('---')
+asyncio.new_event_loop()
+finished.clear()
+asyncio.run(cancel_reader())
+finished.print()
+
+try:
+ asyncio.new_event_loop()
+ finished.clear()
+ asyncio.run(sys_exit())
+except SystemExit:
+ print("SystemExit raised")
+
pFad - Phonifier reborn
Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies:
Alternative Proxy
pFad Proxy
pFad v3 Proxy
pFad v4 Proxy