feat: add Task methods to make _asyncio more similar to cpython

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);
 

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 ""

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())

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

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())

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