asyncio: Properly cancel the main task on exception

greezybacon · greezybacon · commit 4aec5dcdb27a · 2022-11-19T23:46:05.000-06:00
If the main task is interrupted by e.g. a KeyboardInterrupt, then the
main task needs to have the exception injected into it so it will run
the exception handlers and contextmanager __aexit__ methods.
diff --git a/extmod/uasyncio/core.py b/extmod/uasyncio/core.py
@@ -145,9 +145,20 @@ def create_task(coro):
     _task_queue.push(t)
     return t
 
+def run_until_complete(main_task=None):
+    try:
+        _run_until_complete(main_task)
+    except BaseException as exc:
+        try:
+            if main_task:
+                main_task.coro.throw(exc)
+        except StopIteration:
+            # The async task threw the exception
+            pass
+        raise
 
 # Keep scheduling tasks until there are none left to schedule
-def run_until_complete(main_task=None):
+def _run_until_complete(main_task=None):
     global cur_task
     excs_all = (CancelledError, Exception)  # To prevent heap allocation in loop
     excs_stop = (CancelledError, StopIteration)  # To prevent heap allocation in loop
