Skip to content

Commit 335002a

Browse files
committed
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 <[email protected]>
1 parent a41bc5a commit 335002a

File tree

3 files changed

+12
-1
lines changed

3 files changed

+12
-1
lines changed

extmod/moduasyncio.c

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,9 @@ STATIC mp_obj_t task_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
265265
} else if (self->state == TASK_STATE_RUNNING_NOT_WAITED_ON) {
266266
// Allocate the waiting queue.
267267
self->state = task_queue_make_new(&task_queue_type, 0, 0, NULL);
268+
} else if (mp_obj_get_type(self->state) != &task_queue_type) {
269+
// Task has state used for another purpose, so can't also wait on it.
270+
mp_raise_msg(&mp_type_RuntimeError, MP_ERROR_TEXT("can't wait"));
268271
}
269272
return self_in;
270273
}

extmod/uasyncio/core.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,11 @@ def run_until_complete(main_task=None):
195195
if t.state is True:
196196
# "None" indicates that the task is complete and not await'ed on (yet).
197197
t.state = None
198+
elif callable(t.state):
199+
# The task has a callback registered to be called on completion.
200+
t.state(t, er)
201+
t.state = False
202+
waiting = True
198203
else:
199204
# Schedule any other tasks waiting on the completion of this task.
200205
while t.state.peek():

extmod/uasyncio/task.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ class Task:
123123
def __init__(self, coro, globals=None):
124124
self.coro = coro # Coroutine of this Task
125125
self.data = None # General data for queue it is waiting on
126-
self.state = True # None, False, True or a TaskQueue instance
126+
self.state = True # None, False, True, a callable, or a TaskQueue instance
127127
self.ph_key = 0 # Pairing heap
128128
self.ph_child = None # Paring heap
129129
self.ph_child_last = None # Paring heap
@@ -137,6 +137,9 @@ def __iter__(self):
137137
elif self.state is True:
138138
# Allocated head of linked list of Tasks waiting on completion of this task.
139139
self.state = TaskQueue()
140+
elif type(self.state) is not TaskQueue:
141+
# Task has state used for another purpose, so can't also wait on it.
142+
raise RuntimeError("can't wait")
140143
return self
141144

142145
def __next__(self):

0 commit comments

Comments
 (0)