extmod/uasyncio: Allow task state to be a callable.

This implements a form of CPython's "add_done_callback()", but at this
stage it is a hidden feature and only intended to be used internally.

Signed-off-by: Damien George <damien@micropython.org>

1 files changed, 3 insertions, 0 deletions

diff --git a/extmod/moduasyncio.c b/extmod/moduasyncio.c
index dd2d1e747..b0921b6eb 100644
--- a/extmod/moduasyncio.c
+++ b/extmod/moduasyncio.c
@@ -265,6 +265,9 @@ STATIC mp_obj_t task_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
     } 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);
+    } else if (mp_obj_get_type(self->state) != &task_queue_type) {
+        // Task has state used for another purpose, so can't also wait on it.
+        mp_raise_msg(&mp_type_RuntimeError, MP_ERROR_TEXT("can't wait"));
     }
     return self_in;
 }