Add __await__ method to Task class.

Author: t-ikegami

extmod/moduasyncio.c

@@ -29,11 +29,8 @@
 #include "py/pairheap.h"
 #include "py/mphal.h"
 
-#if 0    // causes error for unix build
-#include "shared-bindings/supervisor/__init__.h"
-#else
+// Unix build does not have shared-bindings/supervisor/__init__.h
 extern mp_obj_t supervisor_ticks_ms(void);
-#endif
 
 #if MICROPY_PY_UASYNCIO
 
@@ -232,6 +229,40 @@ STATIC mp_obj_t task_cancel(mp_obj_t self_in) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_cancel_obj, task_cancel);
 
+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);
+    if (TASK_IS_DONE(self)) {
+        // Signal that the completed-task has been await'ed on.
+        self->state = TASK_STATE_DONE_WAS_WAITED_ON;
+    } else if (self->state == TASK_STATE_RUNNING_NOT_WAITED_ON) {
+        // Allocate the waiting queue.
+        self->state = task_queue_make_new(&task_queue_type, 0, 0, NULL);
+    }
+    return self_in;
+}
+
+STATIC mp_obj_t task_iternext(mp_obj_t self_in) {
+    mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
+    if (TASK_IS_DONE(self)) {
+        // Task finished, raise return value to caller so it can continue.
+        nlr_raise(self->data);
+    } else {
+        // Put calling task on waiting queue.
+        mp_obj_t cur_task = mp_obj_dict_get(uasyncio_context, MP_OBJ_NEW_QSTR(MP_QSTR_cur_task));
+        mp_obj_t args[2] = { self->state, cur_task };
+        task_queue_push_sorted(2, args);
+        // Set calling task's data to this task that it waits on, to double-link it.
+        ((mp_obj_task_t *)MP_OBJ_TO_PTR(cur_task))->data = self_in;
+    }
+    return mp_const_none;
+}
+
+STATIC mp_obj_t task_await(mp_obj_t self_in) {
+    return task_getiter(self_in, NULL);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_await_obj, task_await);
+
 STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
     if (dest[0] == MP_OBJ_NULL) {
@@ -250,7 +281,11 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
             dest[1] = self_in;
         } else if (attr == MP_QSTR_ph_key) {
             dest[0] = self->ph_key;
+        } else if (attr == MP_QSTR___await__) {
+            dest[0] = MP_OBJ_FROM_PTR(&task_await_obj);
+            dest[1] = self_in;
         }
+
     } else if (dest[1] != MP_OBJ_NULL) {
         // Store
         if (attr == MP_QSTR_data) {
@@ -263,35 +298,6 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     }
 }
 
-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);
-    if (TASK_IS_DONE(self)) {
-        // Signal that the completed-task has been await'ed on.
-        self->state = TASK_STATE_DONE_WAS_WAITED_ON;
-    } else if (self->state == TASK_STATE_RUNNING_NOT_WAITED_ON) {
-        // Allocate the waiting queue.
-        self->state = task_queue_make_new(&task_queue_type, 0, 0, NULL);
-    }
-    return self_in;
-}
-
-STATIC mp_obj_t task_iternext(mp_obj_t self_in) {
-    mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
-    if (TASK_IS_DONE(self)) {
-        // Task finished, raise return value to caller so it can continue.
-        nlr_raise(self->data);
-    } else {
-        // Put calling task on waiting queue.
-        mp_obj_t cur_task = mp_obj_dict_get(uasyncio_context, MP_OBJ_NEW_QSTR(MP_QSTR_cur_task));
-        mp_obj_t args[2] = { self->state, cur_task };
-        task_queue_push_sorted(2, args);
-        // Set calling task's data to this task that it waits on, to double-link it.
-        ((mp_obj_task_t *)MP_OBJ_TO_PTR(cur_task))->data = self_in;
-    }
-    return mp_const_none;
-}
-
 STATIC const mp_obj_type_t task_type = {
     { &mp_type_type },
     .flags = MP_TYPE_FLAG_EXTENDED,