_poll.pyx revision 781d71db
1"""0MQ polling related functions and classes."""
2
3#
4#    Copyright (c) 2010-2011 Brian E. Granger & Min Ragan-Kelley
5#
6#    This file is part of pyzmq.
7#
8#    pyzmq is free software; you can redistribute it and/or modify it under
9#    the terms of the Lesser GNU General Public License as published by
10#    the Free Software Foundation; either version 3 of the License, or
11#    (at your option) any later version.
12#
13#    pyzmq is distributed in the hope that it will be useful,
14#    but WITHOUT ANY WARRANTY; without even the implied warranty of
15#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
16#    Lesser GNU General Public License for more details.
17#
18#    You should have received a copy of the Lesser GNU General Public License
19#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
20#
21
22#-----------------------------------------------------------------------------
23# Imports
24#-----------------------------------------------------------------------------
25
26from libc.stdlib cimport free, malloc
27
28from libzmq cimport zmq_pollitem_t, ZMQ_VERSION_MAJOR
29from libzmq cimport zmq_poll as zmq_poll_c
30from socket cimport Socket
31
32import sys
33
34from zmq.backend.cython.checkrc cimport _check_rc
35
36#-----------------------------------------------------------------------------
37# Polling related methods
38#-----------------------------------------------------------------------------
39
40# version-independent typecheck for int/long
41if sys.version_info[0] >= 3:
42    int_t = int
43else:
44    int_t = (int,long)
45
46
47def zmq_poll(sockets, long timeout=-1):
48    """zmq_poll(sockets, timeout=-1)
49
50    Poll a set of 0MQ sockets, native file descs. or sockets.
51
52    Parameters
53    ----------
54    sockets : list of tuples of (socket, flags)
55        Each element of this list is a two-tuple containing a socket
56        and a flags. The socket may be a 0MQ socket or any object with
57        a ``fileno()`` method. The flags can be zmq.POLLIN (for detecting
58        for incoming messages), zmq.POLLOUT (for detecting that send is OK)
59        or zmq.POLLIN|zmq.POLLOUT for detecting both.
60    timeout : int
61        The number of milliseconds to poll for. Negative means no timeout.
62    """
63    cdef int rc, i
64    cdef zmq_pollitem_t *pollitems = NULL
65    cdef int nsockets = <int>len(sockets)
66    cdef Socket current_socket
67    
68    if nsockets == 0:
69        return []
70    
71    pollitems = <zmq_pollitem_t *>malloc(nsockets*sizeof(zmq_pollitem_t))
72    if pollitems == NULL:
73        raise MemoryError("Could not allocate poll items")
74        
75    if ZMQ_VERSION_MAJOR < 3:
76        # timeout is us in 2.x, ms in 3.x
77        # expected input is ms (matches 3.x)
78        timeout = 1000*timeout
79    
80    for i in range(nsockets):
81        s, events = sockets[i]
82        if isinstance(s, Socket):
83            pollitems[i].socket = (<Socket>s).handle
84            pollitems[i].events = events
85            pollitems[i].revents = 0
86        elif isinstance(s, int_t):
87            pollitems[i].socket = NULL
88            pollitems[i].fd = s
89            pollitems[i].events = events
90            pollitems[i].revents = 0
91        elif hasattr(s, 'fileno'):
92            try:
93                fileno = int(s.fileno())
94            except:
95                free(pollitems)
96                raise ValueError('fileno() must return a valid integer fd')
97            else:
98                pollitems[i].socket = NULL
99                pollitems[i].fd = fileno
100                pollitems[i].events = events
101                pollitems[i].revents = 0
102        else:
103            free(pollitems)
104            raise TypeError(
105                "Socket must be a 0MQ socket, an integer fd or have "
106                "a fileno() method: %r" % s
107            )
108    
109
110    with nogil:
111        rc = zmq_poll_c(pollitems, nsockets, timeout)
112    
113    if rc < 0:
114        free(pollitems)
115        _check_rc(rc)
116    
117    results = []
118    for i in range(nsockets):
119        revents = pollitems[i].revents
120        # for compatibility with select.poll:
121        # - only return sockets with non-zero status
122        # - return the fd for plain sockets
123        if revents > 0:
124            if pollitems[i].socket != NULL:
125                s = sockets[i][0]
126            else:
127                s = pollitems[i].fd
128            results.append((s, revents))
129
130    free(pollitems)
131    return results
132
133#-----------------------------------------------------------------------------
134# Symbols to export
135#-----------------------------------------------------------------------------
136
137__all__ = [ 'zmq_poll' ]
138