1# coding: utf-8
2"""zmq poll function"""
3
4# Copyright (C) PyZMQ Developers
5# Distributed under the terms of the Modified BSD License.
6
7from ._cffi import C, ffi, zmq_version_info
8
9from .constants import *
10
11from zmq.error import _check_rc
12
13
14def _make_zmq_pollitem(socket, flags):
15    zmq_socket = socket._zmq_socket
16    zmq_pollitem = ffi.new('zmq_pollitem_t*')
17    zmq_pollitem.socket = zmq_socket
18    zmq_pollitem.fd = 0
19    zmq_pollitem.events = flags
20    zmq_pollitem.revents = 0
21    return zmq_pollitem[0]
22
23def _make_zmq_pollitem_fromfd(socket_fd, flags):
24    zmq_pollitem = ffi.new('zmq_pollitem_t*')
25    zmq_pollitem.socket = ffi.NULL
26    zmq_pollitem.fd = socket_fd
27    zmq_pollitem.events = flags
28    zmq_pollitem.revents = 0
29    return zmq_pollitem[0]
30
31def zmq_poll(sockets, timeout):
32    cffi_pollitem_list = []
33    low_level_to_socket_obj = {}
34    for item in sockets:
35        if isinstance(item[0], int):
36            low_level_to_socket_obj[item[0]] = item
37            cffi_pollitem_list.append(_make_zmq_pollitem_fromfd(item[0], item[1]))
38        else:
39            low_level_to_socket_obj[item[0]._zmq_socket] = item
40            cffi_pollitem_list.append(_make_zmq_pollitem(item[0], item[1]))
41    items = ffi.new('zmq_pollitem_t[]', cffi_pollitem_list)
42    list_length = ffi.cast('int', len(cffi_pollitem_list))
43    c_timeout = ffi.cast('long', timeout)
44    rc = C.zmq_poll(items, list_length, c_timeout)
45    _check_rc(rc)
46    result = []
47    for index in range(len(items)):
48        if not items[index].socket == ffi.NULL:
49            if items[index].revents > 0:
50                result.append((low_level_to_socket_obj[items[index].socket][0],
51                            items[index].revents))
52        else:
53            result.append((items[index].fd, items[index].revents))
54    return result
55
56__all__ = ['zmq_poll']
57