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