diff --git a/unix/moduselect.c b/unix/moduselect.c index e74c44be66..f966efa412 100644 --- a/unix/moduselect.c +++ b/unix/moduselect.c @@ -4,7 +4,7 @@ * The MIT License (MIT) * * Copyright (c) 2014 Damien P. George - * Copyright (c) 2015 Paul Sokolovsky + * Copyright (c) 2015-2017 Paul Sokolovsky * * Permission is hereby granted, free of charge, to any person obtaining a copy * of this software and associated documentation files (the "Software"), to deal @@ -55,6 +55,11 @@ typedef struct _mp_obj_poll_t { unsigned short len; struct pollfd *entries; mp_obj_t *obj_map; + short iter_cnt; + short iter_idx; + int flags; + // callee-owned tuple + mp_obj_t ret_tuple; } mp_obj_poll_t; STATIC int get_fd(mp_obj_t fdlike) { @@ -164,9 +169,7 @@ STATIC mp_obj_t poll_modify(mp_obj_t self_in, mp_obj_t obj_in, mp_obj_t eventmas } MP_DEFINE_CONST_FUN_OBJ_3(poll_modify_obj, poll_modify); -/// \method poll([timeout]) -/// Timeout is in milliseconds. -STATIC mp_obj_t poll_poll(size_t n_args, const mp_obj_t *args) { +STATIC int poll_poll_internal(size_t n_args, const mp_obj_t *args) { mp_obj_poll_t *self = MP_OBJ_TO_PTR(args[0]); // work out timeout (it's given already in ms) @@ -184,12 +187,24 @@ STATIC mp_obj_t poll_poll(size_t n_args, const mp_obj_t *args) { } } + self->flags = flags; + int n_ready = poll(self->entries, self->len, timeout); RAISE_ERRNO(n_ready, errno); + return n_ready; +} + +/// \method poll([timeout]) +/// Timeout is in milliseconds. +STATIC mp_obj_t poll_poll(size_t n_args, const mp_obj_t *args) { + int n_ready = poll_poll_internal(n_args, args); + if (n_ready == 0) { return mp_const_empty_tuple; } + mp_obj_poll_t *self = MP_OBJ_TO_PTR(args[0]); + mp_obj_list_t *ret_list = MP_OBJ_TO_PTR(mp_obj_new_list(n_ready, NULL)); int ret_i = 0; struct pollfd *entries = self->entries; @@ -204,7 +219,7 @@ STATIC mp_obj_t poll_poll(size_t n_args, const mp_obj_t *args) { } t->items[1] = MP_OBJ_NEW_SMALL_INT(entries->revents); ret_list->items[ret_i++] = MP_OBJ_FROM_PTR(t); - if (flags & FLAG_ONESHOT) { + if (self->flags & FLAG_ONESHOT) { entries->events = 0; } } @@ -214,17 +229,67 @@ STATIC mp_obj_t poll_poll(size_t n_args, const mp_obj_t *args) { } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(poll_poll_obj, 1, 3, poll_poll); +STATIC mp_obj_t poll_ipoll(size_t n_args, const mp_obj_t *args) { + mp_obj_poll_t *self = MP_OBJ_TO_PTR(args[0]); + + if (self->ret_tuple == MP_OBJ_NULL) { + self->ret_tuple = mp_obj_new_tuple(2, NULL); + } + + int n_ready = poll_poll_internal(n_args, args); + self->iter_cnt = n_ready; + self->iter_idx = 0; + + return args[0]; +} +MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(poll_ipoll_obj, 1, 3, poll_ipoll); + +STATIC mp_obj_t poll_iternext(mp_obj_t self_in) { + mp_obj_poll_t *self = MP_OBJ_TO_PTR(self_in); + + if (self->iter_cnt == 0) { + return MP_OBJ_STOP_ITERATION; + } + + self->iter_cnt--; + + struct pollfd *entries = self->entries; + for (int i = self->iter_idx; i < self->len; i++, entries++) { + if (entries->revents != 0) { + mp_obj_tuple_t *t = MP_OBJ_TO_PTR(self->ret_tuple); + // If there's an object stored, return it, otherwise raw fd + if (self->obj_map && self->obj_map[i] != MP_OBJ_NULL) { + t->items[0] = self->obj_map[i]; + } else { + t->items[0] = MP_OBJ_NEW_SMALL_INT(entries->fd); + } + t->items[1] = MP_OBJ_NEW_SMALL_INT(entries->revents); + if (self->flags & FLAG_ONESHOT) { + entries->events = 0; + } + return MP_OBJ_FROM_PTR(t); + } + } + + assert(!"inconsistent number of poll active entries"); + self->iter_cnt = 0; + return MP_OBJ_STOP_ITERATION; +} + STATIC const mp_rom_map_elem_t poll_locals_dict_table[] = { { MP_ROM_QSTR(MP_QSTR_register), MP_ROM_PTR(&poll_register_obj) }, { MP_ROM_QSTR(MP_QSTR_unregister), MP_ROM_PTR(&poll_unregister_obj) }, { MP_ROM_QSTR(MP_QSTR_modify), MP_ROM_PTR(&poll_modify_obj) }, { MP_ROM_QSTR(MP_QSTR_poll), MP_ROM_PTR(&poll_poll_obj) }, + { MP_ROM_QSTR(MP_QSTR_ipoll), MP_ROM_PTR(&poll_ipoll_obj) }, }; STATIC MP_DEFINE_CONST_DICT(poll_locals_dict, poll_locals_dict_table); STATIC const mp_obj_type_t mp_type_poll = { { &mp_type_type }, .name = MP_QSTR_poll, + .getiter = mp_identity, + .iternext = poll_iternext, .locals_dict = (void*)&poll_locals_dict, }; @@ -239,6 +304,8 @@ STATIC mp_obj_t select_poll(size_t n_args, const mp_obj_t *args) { poll->alloc = alloc; poll->len = 0; poll->obj_map = NULL; + poll->iter_cnt = 0; + poll->ret_tuple = MP_OBJ_NULL; return MP_OBJ_FROM_PTR(poll); } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_select_poll_obj, 0, 1, select_poll);