extmod/modasyncio: Add support for a callback on TaskQueue push.

Allows passing in a callback to `TaskQueue()` that is called when something
is pushed on to the queue.

Signed-off-by: Damien George <damien@micropython.org>
This commit is contained in:
Damien George 2024-06-19 17:27:31 +10:00
parent e9c898cb33
commit 8ac9c8f392
2 changed files with 20 additions and 1 deletions

View File

@ -53,6 +53,9 @@ typedef struct _mp_obj_task_t {
typedef struct _mp_obj_task_queue_t {
mp_obj_base_t base;
mp_obj_task_t *heap;
#if MICROPY_PY_ASYNCIO_TASK_QUEUE_PUSH_CALLBACK
mp_obj_t push_callback;
#endif
} mp_obj_task_queue_t;
static const mp_obj_type_t task_queue_type;
@ -86,9 +89,16 @@ static int task_lt(mp_pairheap_t *n1, mp_pairheap_t *n2) {
static mp_obj_t task_queue_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
(void)args;
mp_arg_check_num(n_args, n_kw, 0, 0, false);
mp_arg_check_num(n_args, n_kw, 0, MICROPY_PY_ASYNCIO_TASK_QUEUE_PUSH_CALLBACK ? 1 : 0, false);
mp_obj_task_queue_t *self = mp_obj_malloc(mp_obj_task_queue_t, type);
self->heap = (mp_obj_task_t *)mp_pairheap_new(task_lt);
#if MICROPY_PY_ASYNCIO_TASK_QUEUE_PUSH_CALLBACK
if (n_args == 1) {
self->push_callback = args[0];
} else {
self->push_callback = MP_OBJ_NULL;
}
#endif
return MP_OBJ_FROM_PTR(self);
}
@ -113,6 +123,11 @@ static mp_obj_t task_queue_push(size_t n_args, const mp_obj_t *args) {
task->ph_key = args[2];
}
self->heap = (mp_obj_task_t *)mp_pairheap_push(task_lt, TASK_PAIRHEAP(self->heap), TASK_PAIRHEAP(task));
#if MICROPY_PY_ASYNCIO_TASK_QUEUE_PUSH_CALLBACK
if (self->push_callback != MP_OBJ_NULL) {
mp_call_function_1(self->push_callback, MP_OBJ_NEW_SMALL_INT(0));
}
#endif
return mp_const_none;
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(task_queue_push_obj, 2, 3, task_queue_push);

View File

@ -1610,6 +1610,10 @@ typedef double mp_float_t;
#define MICROPY_PY_ASYNCIO (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
#endif
#ifndef MICROPY_PY_ASYNCIO_TASK_QUEUE_PUSH_CALLBACK
#define MICROPY_PY_ASYNCIO_TASK_QUEUE_PUSH_CALLBACK (0)
#endif
#ifndef MICROPY_PY_UCTYPES
#define MICROPY_PY_UCTYPES (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
#endif