@@ -182,6 +182,11 @@ MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_select_select_obj, 3, 4, select_select);
182182typedef struct _mp_obj_poll_t {
183183 mp_obj_base_t base ;
184184 mp_map_t poll_map ;
185+ short iter_cnt ;
186+ short iter_idx ;
187+ int flags ;
188+ // callee-owned tuple
189+ mp_obj_t ret_tuple ;
185190} mp_obj_poll_t ;
186191
187192/// \method register(obj[, eventmask])
@@ -219,9 +224,7 @@ STATIC mp_obj_t poll_modify(mp_obj_t self_in, mp_obj_t obj_in, mp_obj_t eventmas
219224}
220225MP_DEFINE_CONST_FUN_OBJ_3 (poll_modify_obj , poll_modify );
221226
222- /// \method poll([timeout])
223- /// Timeout is in milliseconds.
224- STATIC mp_obj_t poll_poll (uint n_args , const mp_obj_t * args ) {
227+ STATIC mp_uint_t poll_poll_internal (uint n_args , const mp_obj_t * args ) {
225228 mp_obj_poll_t * self = args [0 ];
226229
227230 // work out timeout (its given already in ms)
@@ -239,67 +242,130 @@ STATIC mp_obj_t poll_poll(uint n_args, const mp_obj_t *args) {
239242 }
240243 }
241244
245+ self -> flags = flags ;
246+
242247 mp_uint_t start_tick = mp_hal_ticks_ms ();
248+ mp_uint_t n_ready ;
243249 for (;;) {
244250 // poll the objects
245- mp_uint_t n_ready = poll_map_poll (& self -> poll_map , NULL );
246-
251+ n_ready = poll_map_poll (& self -> poll_map , NULL );
247252 if (n_ready > 0 || (timeout != -1 && mp_hal_ticks_ms () - start_tick >= timeout )) {
248- // one or more objects are ready, or we had a timeout
249- mp_obj_list_t * ret_list = mp_obj_new_list (n_ready , NULL );
250- n_ready = 0 ;
251- for (mp_uint_t i = 0 ; i < self -> poll_map .alloc ; ++ i ) {
252- if (!MP_MAP_SLOT_IS_FILLED (& self -> poll_map , i )) {
253- continue ;
254- }
255- poll_obj_t * poll_obj = (poll_obj_t * )self -> poll_map .table [i ].value ;
256- if (poll_obj -> flags_ret != 0 ) {
257- mp_obj_t tuple [2 ] = {poll_obj -> obj , MP_OBJ_NEW_SMALL_INT (poll_obj -> flags_ret )};
258- ret_list -> items [n_ready ++ ] = mp_obj_new_tuple (2 , tuple );
259- if (flags & FLAG_ONESHOT ) {
260- // Don't poll next time, until new event flags will be set explicitly
261- poll_obj -> flags = 0 ;
262- }
263- }
264- }
265- return ret_list ;
253+ break ;
266254 }
267255 MICROPY_EVENT_POLL_HOOK
268256 }
257+
258+ return n_ready ;
259+ }
260+
261+ STATIC mp_obj_t poll_poll (uint n_args , const mp_obj_t * args ) {
262+ mp_obj_poll_t * self = args [0 ];
263+ mp_uint_t n_ready = poll_poll_internal (n_args , args );
264+
265+ // one or more objects are ready, or we had a timeout
266+ mp_obj_list_t * ret_list = mp_obj_new_list (n_ready , NULL );
267+ n_ready = 0 ;
268+ for (mp_uint_t i = 0 ; i < self -> poll_map .alloc ; ++ i ) {
269+ if (!MP_MAP_SLOT_IS_FILLED (& self -> poll_map , i )) {
270+ continue ;
271+ }
272+ poll_obj_t * poll_obj = (poll_obj_t * )self -> poll_map .table [i ].value ;
273+ if (poll_obj -> flags_ret != 0 ) {
274+ mp_obj_t tuple [2 ] = {poll_obj -> obj , MP_OBJ_NEW_SMALL_INT (poll_obj -> flags_ret )};
275+ ret_list -> items [n_ready ++ ] = mp_obj_new_tuple (2 , tuple );
276+ if (self -> flags & FLAG_ONESHOT ) {
277+ // Don't poll next time, until new event flags will be set explicitly
278+ poll_obj -> flags = 0 ;
279+ }
280+ }
281+ }
282+ return ret_list ;
269283}
270284MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN (poll_poll_obj , 1 , 3 , poll_poll );
271285
272- STATIC const mp_map_elem_t poll_locals_dict_table [] = {
273- { MP_OBJ_NEW_QSTR (MP_QSTR_register ), (mp_obj_t )& poll_register_obj },
274- { MP_OBJ_NEW_QSTR (MP_QSTR_unregister ), (mp_obj_t )& poll_unregister_obj },
275- { MP_OBJ_NEW_QSTR (MP_QSTR_modify ), (mp_obj_t )& poll_modify_obj },
276- { MP_OBJ_NEW_QSTR (MP_QSTR_poll ), (mp_obj_t )& poll_poll_obj },
286+ STATIC mp_obj_t poll_ipoll (size_t n_args , const mp_obj_t * args ) {
287+ mp_obj_poll_t * self = MP_OBJ_TO_PTR (args [0 ]);
288+
289+ if (self -> ret_tuple == MP_OBJ_NULL ) {
290+ self -> ret_tuple = mp_obj_new_tuple (2 , NULL );
291+ }
292+
293+ int n_ready = poll_poll_internal (n_args , args );
294+ self -> iter_cnt = n_ready ;
295+ self -> iter_idx = 0 ;
296+
297+ return args [0 ];
298+ }
299+ MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN (poll_ipoll_obj , 1 , 3 , poll_ipoll );
300+
301+ STATIC mp_obj_t poll_iternext (mp_obj_t self_in ) {
302+ mp_obj_poll_t * self = MP_OBJ_TO_PTR (self_in );
303+
304+ if (self -> iter_cnt == 0 ) {
305+ return MP_OBJ_STOP_ITERATION ;
306+ }
307+
308+ self -> iter_cnt -- ;
309+
310+ for (mp_uint_t i = self -> iter_idx ; i < self -> poll_map .alloc ; ++ i ) {
311+ self -> iter_idx ++ ;
312+ if (!MP_MAP_SLOT_IS_FILLED (& self -> poll_map , i )) {
313+ continue ;
314+ }
315+ poll_obj_t * poll_obj = (poll_obj_t * )self -> poll_map .table [i ].value ;
316+ if (poll_obj -> flags_ret != 0 ) {
317+ mp_obj_tuple_t * t = MP_OBJ_TO_PTR (self -> ret_tuple );
318+ t -> items [0 ] = poll_obj -> obj ;
319+ t -> items [1 ] = MP_OBJ_NEW_SMALL_INT (poll_obj -> flags_ret );
320+ if (self -> flags & FLAG_ONESHOT ) {
321+ // Don't poll next time, until new event flags will be set explicitly
322+ poll_obj -> flags = 0 ;
323+ }
324+ return MP_OBJ_FROM_PTR (t );
325+ }
326+ }
327+
328+ assert (!"inconsistent number of poll active entries" );
329+ self -> iter_cnt = 0 ;
330+ return MP_OBJ_STOP_ITERATION ;
331+ }
332+
333+ STATIC const mp_rom_map_elem_t poll_locals_dict_table [] = {
334+ { MP_ROM_QSTR (MP_QSTR_register ), MP_ROM_PTR (& poll_register_obj ) },
335+ { MP_ROM_QSTR (MP_QSTR_unregister ), MP_ROM_PTR (& poll_unregister_obj ) },
336+ { MP_ROM_QSTR (MP_QSTR_modify ), MP_ROM_PTR (& poll_modify_obj ) },
337+ { MP_ROM_QSTR (MP_QSTR_poll ), MP_ROM_PTR (& poll_poll_obj ) },
338+ { MP_ROM_QSTR (MP_QSTR_ipoll ), MP_ROM_PTR (& poll_ipoll_obj ) },
277339};
278340STATIC MP_DEFINE_CONST_DICT (poll_locals_dict , poll_locals_dict_table );
279341
280342STATIC const mp_obj_type_t mp_type_poll = {
281343 { & mp_type_type },
282344 .name = MP_QSTR_poll ,
283- .locals_dict = (mp_obj_t )& poll_locals_dict ,
345+ .getiter = mp_identity_getiter ,
346+ .iternext = poll_iternext ,
347+ .locals_dict = (void * )& poll_locals_dict ,
284348};
285349
286350/// \function poll()
287351STATIC mp_obj_t select_poll (void ) {
288352 mp_obj_poll_t * poll = m_new_obj (mp_obj_poll_t );
289353 poll -> base .type = & mp_type_poll ;
290354 mp_map_init (& poll -> poll_map , 0 );
355+ poll -> iter_cnt = 0 ;
356+ poll -> ret_tuple = MP_OBJ_NULL ;
291357 return poll ;
292358}
293359MP_DEFINE_CONST_FUN_OBJ_0 (mp_select_poll_obj , select_poll );
294360
295- STATIC const mp_map_elem_t mp_module_select_globals_table [] = {
296- { MP_OBJ_NEW_QSTR (MP_QSTR___name__ ), MP_OBJ_NEW_QSTR (MP_QSTR_uselect ) },
297- { MP_OBJ_NEW_QSTR (MP_QSTR_select ), ( mp_obj_t ) & mp_select_select_obj },
298- { MP_OBJ_NEW_QSTR (MP_QSTR_poll ), ( mp_obj_t ) & mp_select_poll_obj },
299- { MP_OBJ_NEW_QSTR (MP_QSTR_POLLIN ), MP_OBJ_NEW_SMALL_INT (MP_STREAM_POLL_RD ) },
300- { MP_OBJ_NEW_QSTR (MP_QSTR_POLLOUT ), MP_OBJ_NEW_SMALL_INT (MP_STREAM_POLL_WR ) },
301- { MP_OBJ_NEW_QSTR (MP_QSTR_POLLERR ), MP_OBJ_NEW_SMALL_INT (MP_STREAM_POLL_ERR ) },
302- { MP_OBJ_NEW_QSTR (MP_QSTR_POLLHUP ), MP_OBJ_NEW_SMALL_INT (MP_STREAM_POLL_HUP ) },
361+ STATIC const mp_rom_map_elem_t mp_module_select_globals_table [] = {
362+ { MP_ROM_QSTR (MP_QSTR___name__ ), MP_ROM_QSTR (MP_QSTR_uselect ) },
363+ { MP_ROM_QSTR (MP_QSTR_select ), MP_ROM_PTR ( & mp_select_select_obj ) },
364+ { MP_ROM_QSTR (MP_QSTR_poll ), MP_ROM_PTR ( & mp_select_poll_obj ) },
365+ { MP_ROM_QSTR (MP_QSTR_POLLIN ), MP_OBJ_NEW_SMALL_INT (MP_STREAM_POLL_RD ) },
366+ { MP_ROM_QSTR (MP_QSTR_POLLOUT ), MP_OBJ_NEW_SMALL_INT (MP_STREAM_POLL_WR ) },
367+ { MP_ROM_QSTR (MP_QSTR_POLLERR ), MP_OBJ_NEW_SMALL_INT (MP_STREAM_POLL_ERR ) },
368+ { MP_ROM_QSTR (MP_QSTR_POLLHUP ), MP_OBJ_NEW_SMALL_INT (MP_STREAM_POLL_HUP ) },
303369};
304370
305371STATIC MP_DEFINE_CONST_DICT (mp_module_select_globals , mp_module_select_globals_table );
0 commit comments