diff --git a/extmod/modasyncio.c b/extmod/modasyncio.c index 4667e3de5332b..bf3f5946df7c5 100644 --- a/extmod/modasyncio.c +++ b/extmod/modasyncio.c @@ -59,6 +59,9 @@ typedef struct _mp_obj_task_queue_t { mp_obj_task_t *heap; } mp_obj_task_queue_t; +MP_DEFINE_EXCEPTION(CancelledError, BaseException) +MP_DEFINE_EXCEPTION(InvalidStateError, Exception) + STATIC const mp_obj_type_t task_queue_type; STATIC const mp_obj_type_t task_type; @@ -202,6 +205,128 @@ STATIC mp_obj_t task_done(mp_obj_t self_in) { } STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_done_obj, task_done); +STATIC mp_obj_t task_add_done_callback(mp_obj_t self_in, mp_obj_t callback) { + assert(mp_obj_is_callable(callback)); + mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); + + if (TASK_IS_DONE(self)) { + // In CPython the callbacks are not immediately called and are instead + // called by the event loop. However, CircuitPython's event loop doesn't + // support `call_soon` to handle callback processing. + // + // Because of this, it's close enough to call the callback immediately. + mp_call_function_1(callback, self_in); + return mp_const_none; + } + + if (self->state != mp_const_true) { + // Tasks SHOULD support more than one callback per CPython but to reduce + // the surface area of this change tasks can currently only support one. + mp_raise_RuntimeError(MP_ERROR_TEXT("Tasks only support one done callback.")); + } + + self->state = callback; + return mp_const_none; +} +STATIC MP_DEFINE_CONST_FUN_OBJ_2(task_add_done_callback_obj, task_add_done_callback); + +STATIC mp_obj_t task_remove_done_callback(mp_obj_t self_in, mp_obj_t callback) { + mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); + + if (callback != self->state) { + return mp_obj_new_int(0); + } + + self->state = mp_const_true; + return mp_obj_new_int(1); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_2(task_remove_done_callback_obj, task_remove_done_callback); + +STATIC mp_obj_t task_get_coro(mp_obj_t self_in) { + mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); + return MP_OBJ_FROM_PTR(self->coro); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_get_coro_obj, task_get_coro); + +STATIC mp_obj_t task_set_exception(mp_obj_t self_in, const mp_obj_t arg) { + mp_raise_msg(&mp_type_RuntimeError, MP_ERROR_TEXT("Task does not support set_exception operation")); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_2(task_set_exception_obj, task_set_exception); + +STATIC mp_obj_t task_exception(mp_obj_t self_in) { + mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); + + if (!TASK_IS_DONE(self)) { + mp_raise_msg(&mp_type_InvalidStateError, MP_ERROR_TEXT("Exception is not set.")); + return NULL; + } + + mp_obj_t data_obj = self->data; + + // If the exception is a cancelled error then we should raise it + if (mp_obj_is_subclass_fast(MP_OBJ_FROM_PTR(mp_obj_get_type(data_obj)), MP_OBJ_FROM_PTR(&mp_type_CancelledError))) { + nlr_raise(data_obj); + } + + // If it's a StopIteration we should should return none + if (mp_obj_is_subclass_fast(MP_OBJ_FROM_PTR(mp_obj_get_type(data_obj)), MP_OBJ_FROM_PTR(&mp_type_StopIteration))) { + return mp_const_none; + } + + if (!mp_obj_is_exception_instance(data_obj)) { + return mp_const_none; + } + + return data_obj; +} +STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_exception_obj, task_exception); + +STATIC mp_obj_t task_set_result(mp_obj_t self_in, const mp_obj_t arg) { + mp_raise_msg(&mp_type_RuntimeError, MP_ERROR_TEXT("Task does not support set_result operation")); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_2(task_set_result_obj, task_set_result); + +STATIC mp_obj_t task_result(mp_obj_t self_in) { + mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); + + if (!TASK_IS_DONE(self)) { + mp_raise_msg(&mp_type_InvalidStateError, MP_ERROR_TEXT("Result is not ready.")); + return NULL; + } + + // If `exception()` returns anything we raise that + mp_obj_t exception_obj = task_exception(self_in); + if (exception_obj != mp_const_none) { + nlr_raise(exception_obj); + } + + mp_obj_t data_obj = self->data; + + // If not StopIteration, bail early + if (!mp_obj_is_subclass_fast(MP_OBJ_FROM_PTR(mp_obj_get_type(data_obj)), MP_OBJ_FROM_PTR(&mp_type_StopIteration))) { + return mp_const_none; + } + + return mp_obj_exception_get_value(data_obj); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_result_obj, task_result); + +STATIC mp_obj_t task_cancelled(mp_obj_t self_in) { + mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); + + if (!TASK_IS_DONE(self)) { + // If task isn't done it can't possibly be cancelled, and would instead + // be considered "cancelling" even if a cancel was requested until it + // has fully completed. + return mp_obj_new_bool(false); + } + + mp_obj_t data_obj = self->data; + + return mp_obj_new_bool(mp_obj_is_subclass_fast(MP_OBJ_FROM_PTR(mp_obj_get_type(data_obj)), MP_OBJ_FROM_PTR(&mp_type_CancelledError))); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_cancelled_obj, task_cancelled); + STATIC mp_obj_t task_cancel(mp_obj_t self_in) { mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); // Check if task is already finished. @@ -276,6 +401,30 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) { } else if (attr == MP_QSTR___await__) { dest[0] = MP_OBJ_FROM_PTR(&task_await_obj); dest[1] = self_in; + } else if (attr == MP_QSTR_add_done_callback) { + dest[0] = MP_OBJ_FROM_PTR(&task_add_done_callback_obj); + dest[1] = self_in; + } else if (attr == MP_QSTR_remove_done_callback) { + dest[0] = MP_OBJ_FROM_PTR(&task_remove_done_callback_obj); + dest[1] = self_in; + } else if (attr == MP_QSTR_get_coro) { + dest[0] = MP_OBJ_FROM_PTR(&task_get_coro_obj); + dest[1] = self_in; + } else if (attr == MP_QSTR_set_result) { + dest[0] = MP_OBJ_FROM_PTR(&task_set_result_obj); + dest[1] = self_in; + } else if (attr == MP_QSTR_result) { + dest[0] = MP_OBJ_FROM_PTR(&task_result_obj); + dest[1] = self_in; + } else if (attr == MP_QSTR_set_exception) { + dest[0] = MP_OBJ_FROM_PTR(&task_set_exception_obj); + dest[1] = self_in; + } else if (attr == MP_QSTR_exception) { + dest[0] = MP_OBJ_FROM_PTR(&task_exception_obj); + dest[1] = self_in; + } else if (attr == MP_QSTR_cancelled) { + dest[0] = MP_OBJ_FROM_PTR(&task_cancelled_obj); + dest[1] = self_in; } } else if (dest[1] != MP_OBJ_NULL) { // Store @@ -289,6 +438,15 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) { } } +STATIC mp_obj_t task_unary_op(mp_unary_op_t op, mp_obj_t o_in) { + switch (op) { + case MP_UNARY_OP_HASH: + return MP_OBJ_NEW_SMALL_INT((mp_uint_t)o_in); + default: + return MP_OBJ_NULL; // op not supported + } +} + STATIC mp_obj_t task_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) { (void)iter_buf; mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); @@ -337,7 +495,8 @@ STATIC MP_DEFINE_CONST_OBJ_TYPE( MP_TYPE_FLAG_ITER_IS_CUSTOM, make_new, task_make_new, attr, task_attr, - iter, &task_getiter_iternext + iter, &task_getiter_iternext, + unary_op, task_unary_op ); /******************************************************************************/ @@ -347,6 +506,8 @@ STATIC const mp_rom_map_elem_t mp_module_asyncio_globals_table[] = { { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR__asyncio) }, { MP_ROM_QSTR(MP_QSTR_TaskQueue), MP_ROM_PTR(&task_queue_type) }, { MP_ROM_QSTR(MP_QSTR_Task), MP_ROM_PTR(&task_type) }, + { MP_ROM_QSTR(MP_QSTR_CancelledError), MP_ROM_PTR(&mp_type_CancelledError) }, + { MP_ROM_QSTR(MP_QSTR_InvalidStateError), MP_ROM_PTR(&mp_type_InvalidStateError) }, }; STATIC MP_DEFINE_CONST_DICT(mp_module_asyncio_globals, mp_module_asyncio_globals_table); diff --git a/locale/circuitpython.pot b/locale/circuitpython.pot index f2f21f82d51ae..e052a3c28abba 100644 --- a/locale/circuitpython.pot +++ b/locale/circuitpython.pot @@ -958,6 +958,10 @@ msgstr "" msgid "Error: Failure to bind" msgstr "" +#: extmod/modasyncio.c +msgid "Exception is not set." +msgstr "" + #: shared-bindings/alarm/__init__.c msgid "Expected a kind of %q" msgstr "" @@ -1858,6 +1862,10 @@ msgstr "" msgid "Requested resource not found" msgstr "" +#: extmod/modasyncio.c +msgid "Result is not ready." +msgstr "" + #: ports/atmel-samd/common-hal/audioio/AudioOut.c msgid "Right channel unsupported" msgstr "" @@ -1963,6 +1971,18 @@ msgstr "" msgid "System entry must be gnss.SatelliteSystem" msgstr "" +#: extmod/modasyncio.c +msgid "Task does not support set_exception operation" +msgstr "" + +#: extmod/modasyncio.c +msgid "Task does not support set_result operation" +msgstr "" + +#: extmod/modasyncio.c +msgid "Tasks only support one done callback." +msgstr "" + #: ports/stm/common-hal/microcontroller/Processor.c msgid "Temperature read timed out" msgstr "" diff --git a/tests/extmod/asyncio_task_add_done_callback.py b/tests/extmod/asyncio_task_add_done_callback.py new file mode 100644 index 0000000000000..899479c9d3070 --- /dev/null +++ b/tests/extmod/asyncio_task_add_done_callback.py @@ -0,0 +1,50 @@ +# Test the Task.done() method + +try: + import asyncio +except ImportError: + print("SKIP") + raise SystemExit + + +async def task(t, exc=None): + if t >= 0: + await asyncio.sleep(t) + if exc: + raise exc + + +async def main(): + # Tasks that aren't done only execute done callback after finishing + print("=" * 10) + t = asyncio.create_task(task(-1)) + t.add_done_callback(lambda: print("done")) + print("Waiting for task to complete") + await asyncio.sleep(0) + print("Task has completed") + + # Task that are done run the callback immediately + print("=" * 10) + t = asyncio.create_task(task(-1)) + await asyncio.sleep(0) + print("Task has completed") + t.add_done_callback(lambda: print("done")) + print("Callback Added") + + # Task that starts, runs and finishes without an exception should return None + print("=" * 10) + t = asyncio.create_task(task(0.01)) + t.add_done_callback(lambda: print("done")) + try: + t.add_done_callback(lambda: print("done")) + except RuntimeError as e: + print("Second call to add_done_callback emits error:", repr(e)) + + # Task that raises immediately should still run done callback + print("=" * 10) + t = asyncio.create_task(task(-1, ValueError)) + t.add_done_callback(lambda: print("done")) + await asyncio.sleep(0) + + +asyncio.run(main()) diff --git a/tests/extmod/asyncio_task_add_done_callback.py.exp b/tests/extmod/asyncio_task_add_done_callback.py.exp new file mode 100644 index 0000000000000..fa20e94d6845f --- /dev/null +++ b/tests/extmod/asyncio_task_add_done_callback.py.exp @@ -0,0 +1,14 @@ +========== +Waiting for task to complete +done +Task has completed +========== +Task has completed +done +Callback added +========== +Second call to add_done_callback emits error: Tasks only support one done callback. +========== +Waiting for task to complete +done +Exception handled diff --git a/tests/extmod/asyncio_task_cancelled.py b/tests/extmod/asyncio_task_cancelled.py new file mode 100644 index 0000000000000..763a0fc987e42 --- /dev/null +++ b/tests/extmod/asyncio_task_cancelled.py @@ -0,0 +1,54 @@ +# Test cancelling a task + +try: + import asyncio +except ImportError: + print("SKIP") + raise SystemExit + + +async def task(t): + await asyncio.sleep(t) + + +async def main(): + # Cancel task immediately doesn't mark the task as cancelled + print("=" * 10) + t = asyncio.create_task(task(2)) + t.cancel() + print("Expecting task to not be cancelled because it is not done:", t.cancelled()) + + # Cancel task immediately and wait for cancellation to complete + print("=" * 10) + t = asyncio.create_task(task(2)) + t.cancel() + await asyncio.sleep(0) + print("Expecting Task to be Cancelled:", t.cancelled()) + + # Cancel task and wait for cancellation to complete + print("=" * 10) + t = asyncio.create_task(task(2)) + await asyncio.sleep(0.01) + t.cancel() + await asyncio.sleep(0) + print("Expecting Task to be Cancelled:", t.cancelled()) + + # Cancel task multiple times after it has started + print("=" * 10) + t = asyncio.create_task(task(2)) + await asyncio.sleep(0.01) + for _ in range(4): + t.cancel() + await asyncio.sleep(0.01) + + print("Expecting Task to be Cancelled:", t.cancelled()) + + # Cancel task after it has finished + print("=" * 10) + t = asyncio.create_task(task(0.01)) + await asyncio.sleep(0.05) + t.cancel() + print("Expecting task to not be Cancelled:", t.cancelled()) + + +asyncio.run(main()) diff --git a/tests/extmod/asyncio_task_cancelled.py.exp b/tests/extmod/asyncio_task_cancelled.py.exp new file mode 100644 index 0000000000000..7539d2bc1d9bb --- /dev/null +++ b/tests/extmod/asyncio_task_cancelled.py.exp @@ -0,0 +1,10 @@ +========== +Expecting task to not be cancelled because it is not done: False +========== +Expecting Task to be Cancelled: True +========== +Expecting Task to be Cancelled: True +========== +Expecting Task to be Cancelled: True +========== +Expecting task to not be Cancelled: False \ No newline at end of file diff --git a/tests/extmod/asyncio_task_exception.py b/tests/extmod/asyncio_task_exception.py new file mode 100644 index 0000000000000..7ae4e588403bf --- /dev/null +++ b/tests/extmod/asyncio_task_exception.py @@ -0,0 +1,57 @@ +# Test the Task.done() method + +try: + import asyncio +except ImportError: + print("SKIP") + raise SystemExit + + +async def task(t, exc=None): + if t >= 0: + await asyncio.sleep(t) + if exc: + raise exc + + +async def main(): + # Task that is not done yet raises an InvalidStateError + print("=" * 10) + t = asyncio.create_task(task(1)) + await asyncio.sleep(0) + try: + t.exception() + assert False, "Should not get here" + except Exception as e: + print("Tasks that aren't done yet raise an InvalidStateError:", repr(e)) + + # Task that is cancelled raises CancelledError + print("=" * 10) + t = asyncio.create_task(task(1)) + t.cancel() + await asyncio.sleep(0) + try: + print(repr(t.exception())) + print(t.cancelled()) + assert False, "Should not get here" + except asyncio.CancelledError as e: + print("Cancelled tasks cannot retrieve exception:", repr(e)) + + # Task that starts, runs and finishes without an exception should return None + print("=" * 10) + t = asyncio.create_task(task(0.01)) + await t + print("None when no exception:", t.exception()) + + # Task that raises immediately should return that exception + print("=" * 10) + t = asyncio.create_task(task(-1, ValueError)) + try: + await t + assert False, "Should not get here" + except ValueError as e: + pass + print("Returned Exception:", repr(t.exception())) + + +asyncio.run(main()) diff --git a/tests/extmod/asyncio_task_exception.py.exp b/tests/extmod/asyncio_task_exception.py.exp new file mode 100644 index 0000000000000..682c240f937ee --- /dev/null +++ b/tests/extmod/asyncio_task_exception.py.exp @@ -0,0 +1,8 @@ +========== +Tasks that aren't done yet raise an InvalidStateError: InvalidStateError('Exception is not set.',) +========== +Cancelled tasks cannot retrieve exception: CancelledError() +========== +None when no exception: None +========== +Returned Exception: ValueError() \ No newline at end of file diff --git a/tests/extmod/asyncio_task_remove_done_callback.py b/tests/extmod/asyncio_task_remove_done_callback.py new file mode 100644 index 0000000000000..c3429ec3b5717 --- /dev/null +++ b/tests/extmod/asyncio_task_remove_done_callback.py @@ -0,0 +1,59 @@ +# Test the Task.done() method + +try: + import asyncio +except ImportError: + print("SKIP") + raise SystemExit + + +async def task(t, exc=None): + if t >= 0: + await asyncio.sleep(t) + if exc: + raise exc + + +def done_callback(): + print("done") + + +def done_callback_2(): + print("done 2") + + +async def main(): + # Removing a callback returns 0 when no callbacks have been set + print("=" * 10) + t = asyncio.create_task(task(1)) + print("Returns 0 when no done callback has been set:", t.remove_done_callback(done_callback)) + + # Done callback removal only works once + print("=" * 10) + t = asyncio.create_task(task(1)) + t.add_done_callback(done_callback) + print( + "Returns 1 when a callback matches and is removed:", t.remove_done_callback(done_callback) + ) + print( + "Returns 0 on second attempt to remove the callback:", + t.remove_done_callback(done_callback), + ) + + # Only removes done callback when match + print("=" * 10) + t = asyncio.create_task(task(0.01)) + t.add_done_callback(done_callback) + print("Returns 0 when done callbacks don't match:", t.remove_done_callback(done_callback_2)) + + # A removed done callback does not execute + print("=" * 10) + t = asyncio.create_task(task(-1)) + t.add_done_callback(done_callback) + t.remove_done_callback(done_callback) + print("Waiting for task to complete") + await t + print("Task completed") + + +asyncio.run(main()) diff --git a/tests/extmod/asyncio_task_remove_done_callback.py.exp b/tests/extmod/asyncio_task_remove_done_callback.py.exp new file mode 100644 index 0000000000000..50c0b63dd0675 --- /dev/null +++ b/tests/extmod/asyncio_task_remove_done_callback.py.exp @@ -0,0 +1,10 @@ +========== +Returns 0 when no done callback has been set: 0 +========== +Returns 1 when a callback matches and is removed: 1 +Returns 0 on second attempt to remove the callback: 0 +========== +Returns 0 when done callbacks don't match: 0 +========== +Waiting for task to complete +Task completed diff --git a/tests/extmod/asyncio_task_result.py b/tests/extmod/asyncio_task_result.py new file mode 100644 index 0000000000000..a389df18c2007 --- /dev/null +++ b/tests/extmod/asyncio_task_result.py @@ -0,0 +1,69 @@ +# Test the Task.done() method + +try: + import asyncio +except ImportError: + print("SKIP") + raise SystemExit + + +async def task(t, exc=None, ret=None): + if t >= 0: + await asyncio.sleep(t) + if exc: + raise exc + return ret + + +async def main(): + # Task that is not done yet raises an InvalidStateError + print("=" * 10) + t = asyncio.create_task(task(1)) + await asyncio.sleep(0) + try: + t.result() + assert False, "Should not get here" + except Exception as e: + print("InvalidStateError if still running:", repr(e)) + + # Task that is cancelled raises CancelledError + print("=" * 10) + t = asyncio.create_task(task(1)) + t.cancel() + await asyncio.sleep(0) + try: + t.result() + assert False, "Should not get here" + except asyncio.CancelledError as e: + print("CancelledError when retrieving result from cancelled task:", repr(e)) + + # Task that raises immediately should raise that exception when calling result + print("=" * 10) + t = asyncio.create_task(task(-1, ValueError)) + try: + await t + assert False, "Should not get here" + except ValueError as e: + pass + + try: + t.result() + assert False, "Should not get here" + except ValueError as e: + print("Error raised when result is attempted on task with error:", repr(e)) + + # Task that starts, runs and finishes without an exception or value should return None + print("=" * 10) + t = asyncio.create_task(task(0.01)) + await t + print("Empty Result should be None:", t.result()) + assert t.result() is None + + # Task that starts, runs and finishes without exception should return result + print("=" * 10) + t = asyncio.create_task(task(0.01, None, "hello world")) + await t + print("Happy path, result is returned:", t.result()) + + +asyncio.run(main()) diff --git a/tests/extmod/asyncio_task_result.py.exp b/tests/extmod/asyncio_task_result.py.exp new file mode 100644 index 0000000000000..0773485be60a6 --- /dev/null +++ b/tests/extmod/asyncio_task_result.py.exp @@ -0,0 +1,10 @@ +========== +InvalidStateError if still running: InvalidStateError('Result is not ready.',) +========== +CancelledError when retrieving result from cancelled task: CancelledError() +========== +Error raised when result is attempted on task with error: ValueError() +========== +Empty Result should be None: None +========== +Happy path, result is returned: hello world