Add __await__ method to Task class.

This commit is contained in:
Tsutomu IKEGAMI 2022-02-22 17:11:11 +09:00
parent 0929353e5c
commit 6b7ad72f8f
1 changed files with 41 additions and 35 deletions

View File

@ -29,11 +29,8 @@
#include "py/pairheap.h" #include "py/pairheap.h"
#include "py/mphal.h" #include "py/mphal.h"
#if 0 // causes error for unix build // Unix build does not have shared-bindings/supervisor/__init__.h
#include "shared-bindings/supervisor/__init__.h"
#else
extern mp_obj_t supervisor_ticks_ms(void); extern mp_obj_t supervisor_ticks_ms(void);
#endif
#if MICROPY_PY_UASYNCIO #if MICROPY_PY_UASYNCIO
@ -232,37 +229,6 @@ 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_DEFINE_CONST_FUN_OBJ_1(task_cancel_obj, task_cancel);
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) {
// Load
if (attr == MP_QSTR_coro) {
dest[0] = self->coro;
} else if (attr == MP_QSTR_data) {
dest[0] = self->data;
} else if (attr == MP_QSTR_state) {
dest[0] = self->state;
} 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;
} else if (attr == MP_QSTR_ph_key) {
dest[0] = self->ph_key;
}
} else if (dest[1] != MP_OBJ_NULL) {
// Store
if (attr == MP_QSTR_data) {
self->data = dest[1];
dest[0] = MP_OBJ_NULL;
} else if (attr == MP_QSTR_state) {
self->state = dest[1];
dest[0] = MP_OBJ_NULL;
}
}
}
STATIC mp_obj_t task_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) { STATIC mp_obj_t task_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
(void)iter_buf; (void)iter_buf;
mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in); mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
@ -292,6 +258,46 @@ STATIC mp_obj_t task_iternext(mp_obj_t self_in) {
return mp_const_none; 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) {
// Load
if (attr == MP_QSTR_coro) {
dest[0] = self->coro;
} else if (attr == MP_QSTR_data) {
dest[0] = self->data;
} else if (attr == MP_QSTR_state) {
dest[0] = self->state;
} 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;
} 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) {
self->data = dest[1];
dest[0] = MP_OBJ_NULL;
} else if (attr == MP_QSTR_state) {
self->state = dest[1];
dest[0] = MP_OBJ_NULL;
}
}
}
STATIC const mp_obj_type_t task_type = { STATIC const mp_obj_type_t task_type = {
{ &mp_type_type }, { &mp_type_type },
.flags = MP_TYPE_FLAG_EXTENDED, .flags = MP_TYPE_FLAG_EXTENDED,