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>

1 files changed, 16 insertions, 1 deletions

diff --git a/extmod/modasyncio.c b/extmod/modasyncio.c
index 4f82370d2..6161500bf 100644
--- a/extmod/modasyncio.c
+++ b/extmod/modasyncio.c
@@ -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);