extmod/uasyncio: Add Task.done() method.

dpgeorge · dpgeorge · commit 309dfe39e07d · 2020-12-02T12:07:06.000+11:00
This is added because task.coro==None is no longer the way to detect if a
task is finished.  Providing a (CPython compatible) function for this
allows the implementation to be abstracted away.

Signed-off-by: Damien George &lt;damien@micropython.org&gt;
diff --git a/extmod/moduasyncio.c b/extmod/moduasyncio.c
@@ -167,6 +167,12 @@ STATIC mp_obj_t task_make_new(const mp_obj_type_t *type, size_t n_args, size_t n
     return MP_OBJ_FROM_PTR(self);
 }
 
+STATIC mp_obj_t task_done(mp_obj_t self_in) {
+    mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
+    return mp_obj_new_bool(TASK_IS_DONE(self));
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_done_obj, task_done);
+
 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.
@@ -242,6 +248,9 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
             if (self->waiting != mp_const_none && self->waiting != mp_const_false) {
                 dest[0] = self->waiting;
             }
+        } else if (attr == MP_QSTR_done) {
+            dest[0] = MP_OBJ_FROM_PTR(&task_done_obj);
+            dest[1] = self_in;
         } else if (attr == MP_QSTR_cancel) {
             dest[0] = MP_OBJ_FROM_PTR(&task_cancel_obj);
             dest[1] = self_in;
diff --git a/extmod/uasyncio/task.py b/extmod/uasyncio/task.py
@@ -148,6 +148,9 @@ def __next__(self):
             # Set calling task's data to this task that it waits on, to double-link it.
             core.cur_task.data = self
 
+    def done(self):
+        return self.coro is self
+
     def cancel(self):
         # Check if task is already finished.
         if self.coro is self:
diff --git a/tests/extmod/uasyncio_task_done.py b/tests/extmod/uasyncio_task_done.py
@@ -0,0 +1,66 @@
+# Test the Task.done() method
+
+try:
+    import uasyncio as asyncio
+except ImportError:
+    try:
+        import asyncio
+    except ImportError:
+        print("SKIP")
+        raise SystemExit
+
+
+async def task(t, exc=None):
+    print("task start")
+    if t >= 0:
+        await asyncio.sleep(t)
+    if exc:
+        raise exc
+    print("task done")
+
+
+async def main():
+    # Task that finishes immediately.
+    print("=" * 10)
+    t = asyncio.create_task(task(-1))
+    print(t.done())
+    await asyncio.sleep(0)
+    print(t.done())
+    await t
+    print(t.done())
+
+    # Task that starts, runs and finishes.
+    print("=" * 10)
+    t = asyncio.create_task(task(0.01))
+    print(t.done())
+    await asyncio.sleep(0)
+    print(t.done())
+    await t
+    print(t.done())
+
+    # Task that raises immediately.
+    print("=" * 10)
+    t = asyncio.create_task(task(-1, ValueError))
+    print(t.done())
+    await asyncio.sleep(0)
+    print(t.done())
+    try:
+        await t
+    except ValueError as er:
+        print(repr(er))
+    print(t.done())
+
+    # Task that raises after a delay.
+    print("=" * 10)
+    t = asyncio.create_task(task(0.01, ValueError))
+    print(t.done())
+    await asyncio.sleep(0)
+    print(t.done())
+    try:
+        await t
+    except ValueError as er:
+        print(repr(er))
+    print(t.done())
+
+
+asyncio.run(main())
diff --git a/tests/extmod/uasyncio_task_done.py.exp b/tests/extmod/uasyncio_task_done.py.exp
@@ -0,0 +1,24 @@
+==========
+False
+task start
+task done
+True
+True
+==========
+False
+task start
+False
+task done
+True
+==========
+False
+task start
+True
+ValueError()
+True
+==========
+False
+task start
+False
+ValueError()
+True
