1 files changed, 72 insertions, 5 deletions

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);