From 309dfe39e07d3c3c8ba873ed09116cea9f6f52c0 Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Mon, 30 Nov 2020 17:38:30 +1100
Subject: [PATCH] extmod/uasyncio: Add Task.done() method.

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 <damien@micropython.org>
---
 extmod/moduasyncio.c                   |  9 ++++
 extmod/uasyncio/task.py                |  3 ++
 tests/extmod/uasyncio_task_done.py     | 66 ++++++++++++++++++++++++++
 tests/extmod/uasyncio_task_done.py.exp | 24 ++++++++++
 4 files changed, 102 insertions(+)
 create mode 100644 tests/extmod/uasyncio_task_done.py
 create mode 100644 tests/extmod/uasyncio_task_done.py.exp

diff --git a/extmod/moduasyncio.c b/extmod/moduasyncio.c
index e8822c0697..fe0f748cac 100644
--- 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
index 2420ab7193..68ddf496f0 100644
--- a/extmod/uasyncio/task.py
+++ b/extmod/uasyncio/task.py
@@ -148,6 +148,9 @@ class Task:
             # 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
new file mode 100644
index 0000000000..2700da8c34
--- /dev/null
+++ 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
new file mode 100644
index 0000000000..ddda04c5ec
--- /dev/null
+++ 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