Add interface to poll() system call (SF patch #100852)

Modules/selectmodule.c

@@ -24,6 +24,9 @@ redistribution of this file, and for a DISCLAIMER OF ALL WARRANTIES.
 #ifdef HAVE_LIMITS_H
 #include <limits.h>
 #endif
+#ifdef HAVE_POLL_H
+#include <poll.h>
+#endif
 
 #ifdef __sgi
 /* This is missing from unistd.h */
@@ -299,6 +302,284 @@ select_select(PyObject *self, PyObject *args)
 	return ret;
 }
 
+#ifdef HAVE_POLL
+/* 
+ * poll() support
+ */
+
+typedef struct {
+	PyObject_HEAD
+	PyObject *dict;
+	int ufd_uptodate; 
+	int ufd_len;
+        struct pollfd *ufds;
+} pollObject;
+
+staticforward PyTypeObject poll_Type;
+
+/* Update the malloc'ed array of pollfds to match the dictionary 
+   contained within a pollObject.  Return 1 on success, 0 on an error.
+*/
+
+static int
+update_ufd_array(pollObject *self)
+{
+	int i, j, pos;
+	PyObject *key, *value;
+
+	self->ufd_len = PyDict_Size(self->dict);
+	PyMem_Resize(self->ufds, struct pollfd, self->ufd_len);
+	if (self->ufds == NULL) {
+		PyErr_NoMemory();
+		return 0;
+	}
+
+	i = pos = 0;
+	while ((j = PyDict_Next(self->dict, &pos, &key, &value))) {
+		self->ufds[i].fd = PyInt_AsLong(key);
+		self->ufds[i].events = PyInt_AsLong(value);
+		i++;
+	}
+	self->ufd_uptodate = 1;
+	return 1;
+}
+
+static char poll_register_doc[] =
+"register(fd [, eventmask] ) -> None\n\n\
+Register a file descriptor with the polling object.\n\
+fd -- either an integer, or an object with a fileno() method returning an int.\n\
+events -- an optional bitmask describing the type of events to check for";
+
+static PyObject *
+poll_register(pollObject *self, PyObject *args) 
+{
+	PyObject *o, *key, *value;
+	int fd, events = POLLIN | POLLPRI | POLLOUT;
+
+	if (!PyArg_ParseTuple(args, "O|i", &o, &events)) {
+		return NULL;
+	}
+  
+	fd = PyObject_AsFileDescriptor(o);
+	if (fd == -1) return NULL;
+
+	/* Add entry to the internal dictionary: the key is the 
+	   file descriptor, and the value is the event mask. */
+	if ( (NULL == (key = PyInt_FromLong(fd))) ||
+	     (NULL == (value = PyInt_FromLong(events))) ||
+	     (PyDict_SetItem(self->dict, key, value)) == -1) {
+		return NULL;
+	}
+	self->ufd_uptodate = 0;
+		       
+	Py_INCREF(Py_None);
+	return Py_None;
+}
+
+static char poll_unregister_doc[] =
+"unregister(fd) -> None\n\n\
+Remove a file descriptor being tracked by the polling object.";
+
+static PyObject *
+poll_unregister(pollObject *self, PyObject *args) 
+{
+	PyObject *o, *key;
+	int fd;
+
+	if (!PyArg_ParseTuple(args, "O", &o)) {
+		return NULL;
+	}
+  
+	fd = PyObject_AsFileDescriptor( o );
+	if (fd == -1) 
+		return NULL;
+
+	/* Check whether the fd is already in the array */
+	key = PyInt_FromLong(fd);
+	if (key == NULL) 
+		return NULL;
+
+	if (PyDict_DelItem(self->dict, key) == -1) {
+		Py_DECREF(key);
+		/* This will simply raise the KeyError set by PyDict_DelItem
+		   if the file descriptor isn't registered. */
+		return NULL;
+	}
+
+	Py_DECREF(key);
+	self->ufd_uptodate = 0;
+
+	Py_INCREF(Py_None);
+	return Py_None;
+}
+
+static char poll_poll_doc[] =
+"poll( [timeout] ) -> list of (fd, event) 2-tuples\n\n\
+Polls the set of registered file descriptors, returning a list containing \n\
+any descriptors that have events or errors to report.";
+
+static PyObject *
+poll_poll(pollObject *self, PyObject *args) 
+{
+	PyObject *result_list = NULL, *tout = NULL;
+	int timeout = 0, poll_result, i, j;
+	PyObject *value = NULL, *num = NULL;
+
+	if (!PyArg_ParseTuple(args, "|O", &tout)) {
+		return NULL;
+	}
+
+	/* Check values for timeout */
+	if (tout == NULL || tout == Py_None)
+		timeout = -1;
+	else if (!PyArg_Parse(tout, "i", &timeout)) {
+		PyErr_SetString(PyExc_TypeError,
+				"timeout must be an integer or None");
+		return NULL;
+	}
+
+	/* Ensure the ufd array is up to date */
+	if (!self->ufd_uptodate) 
+		if (update_ufd_array(self) == 0)
+			return NULL;
+
+	/* call poll() */
+	Py_BEGIN_ALLOW_THREADS;
+	poll_result = poll(self->ufds, self->ufd_len, timeout);
+	Py_END_ALLOW_THREADS;
+ 
+	if (poll_result < 0) {
+		PyErr_SetFromErrno(SelectError);
+		return NULL;
+	} 
+       
+	/* build the result list */
+  
+	result_list = PyList_New(poll_result);
+	if (!result_list) 
+		return NULL;
+	else {
+		for (i = 0, j = 0; j < poll_result; j++) {
+ 			/* skip to the next fired descriptor */
+ 			while (!self->ufds[i].revents) {
+ 				i++;
+ 			}
+			/* if we hit a NULL return, set value to NULL
+			   and break out of loop; code at end will
+			   clean up result_list */
+			value = PyTuple_New(2);
+			if (value == NULL)
+				goto error;
+			num = PyInt_FromLong(self->ufds[i].fd);
+			if (num == NULL) {
+				Py_DECREF(value);
+				goto error;
+			}
+			PyTuple_SET_ITEM(value, 0, num);
+
+			num = PyInt_FromLong(self->ufds[i].revents);
+			if (num == NULL) {
+				Py_DECREF(value);
+				goto error;
+			}
+			PyTuple_SET_ITEM(value, 1, num);
+ 			if ((PyList_SetItem(result_list, j, value)) == -1) {
+				Py_DECREF(value);
+				goto error;
+ 			}
+ 			i++;
+ 		}
+ 	}
+ 	return result_list;
+
+  error:
+	Py_DECREF(result_list);
+	return NULL;
+}
+
+static PyMethodDef poll_methods[] = {
+	{"register",	(PyCFunction)poll_register,	
+	 METH_VARARGS,  poll_register_doc},
+	{"unregister",	(PyCFunction)poll_unregister,	
+	 METH_VARARGS,  poll_unregister_doc},
+	{"poll",	(PyCFunction)poll_poll,	
+	 METH_VARARGS,  poll_poll_doc},
+	{NULL,		NULL}		/* sentinel */
+};
+
+static pollObject *
+newPollObject()
+{
+        pollObject *self;
+	self = PyObject_New(pollObject, &poll_Type);
+	if (self == NULL)
+		return NULL;
+	/* ufd_uptodate is a Boolean, denoting whether the 
+	   array pointed to by ufds matches the contents of the dictionary. */
+	self->ufd_uptodate = 0;
+	self->ufds = NULL;
+	self->dict = PyDict_New();
+	if (self->dict == NULL) {
+		Py_DECREF(self);
+		return NULL;
+	}
+	return self;
+}
+
+static void
+poll_dealloc(pollObject *self)
+{
+	if (self->ufds != NULL)
+		PyMem_DEL(self->ufds);
+	Py_XDECREF(self->dict);
+  	PyObject_Del(self);
+}
+
+static PyObject *
+poll_getattr(pollObject *self, char *name)
+{
+	return Py_FindMethod(poll_methods, (PyObject *)self, name);
+}
+
+statichere PyTypeObject poll_Type = {
+	/* The ob_type field must be initialized in the module init function
+	 * to be portable to Windows without using C++. */
+	PyObject_HEAD_INIT(NULL)
+	0,			/*ob_size*/
+	"poll",			/*tp_name*/
+	sizeof(pollObject),	/*tp_basicsize*/
+	0,			/*tp_itemsize*/
+	/* methods */
+	(destructor)poll_dealloc, /*tp_dealloc*/
+	0,			/*tp_print*/
+	(getattrfunc)poll_getattr, /*tp_getattr*/
+	0,                      /*tp_setattr*/
+	0,			/*tp_compare*/
+	0,			/*tp_repr*/
+	0,			/*tp_as_number*/
+	0,			/*tp_as_sequence*/
+	0,			/*tp_as_mapping*/
+	0,			/*tp_hash*/
+};
+
+static char poll_doc[] = 
+"Returns a polling object, which supports registering and\n\
+unregistering file descriptors, and then polling them for I/O events.";
+
+static PyObject *
+select_poll(PyObject *self, PyObject *args)
+{
+	pollObject *rv;
+	
+	if (!PyArg_ParseTuple(args, ":poll"))
+		return NULL;
+	rv = newPollObject();
+	if ( rv == NULL )
+		return NULL;
+	return (PyObject *)rv;
+}
+#endif /* HAVE_POLL */
+
 static char select_doc[] =
 "select(rlist, wlist, xlist[, timeout]) -> (rlist, wlist, xlist)\n\
 \n\
@@ -322,9 +603,11 @@ that are ready.\n\
 *** IMPORTANT NOTICE ***\n\
 On Windows, only sockets are supported; on Unix, all file descriptors.";
 
-
 static PyMethodDef select_methods[] = {
-    {"select",	select_select, 1, select_doc},
+    {"select",	select_select, METH_VARARGS, select_doc},
+#ifdef HAVE_POLL
+    {"poll",    select_poll,   METH_VARARGS, poll_doc},
+#endif /* HAVE_POLL */
     {0,  	0},			     /* sentinel */
 };
 
@@ -334,6 +617,25 @@ static char module_doc[] =
 *** IMPORTANT NOTICE ***\n\
 On Windows, only sockets are supported; on Unix, all file descriptors.";
 
+/*
+ * Convenience routine to export an integer value.
+ * For simplicity, errors (which are unlikely anyway) are ignored.
+ */
+
+static void
+insint(PyObject *d, char *name, int value)
+{
+	PyObject *v = PyInt_FromLong((long) value);
+	if (v == NULL) {
+		/* Don't bother reporting this error */
+		PyErr_Clear();
+	}
+	else {
+		PyDict_SetItemString(d, name, v);
+		Py_DECREF(v);
+	}
+}
+
 DL_EXPORT(void)
 initselect(void)
 {
@@ -342,4 +644,19 @@ initselect(void)
 	d = PyModule_GetDict(m);
 	SelectError = PyErr_NewException("select.error", NULL, NULL);
 	PyDict_SetItemString(d, "error", SelectError);
+#ifdef HAVE_POLL
+	poll_Type.ob_type = &PyType_Type;
+	insint(d, "POLLIN", POLLIN);
+	insint(d, "POLLPRI", POLLPRI);
+	insint(d, "POLLOUT", POLLOUT);
+	insint(d, "POLLERR", POLLERR);
+	insint(d, "POLLHUP", POLLHUP);
+	insint(d, "POLLNVAL", POLLNVAL);
+
+	insint(d, "POLLRDNORM", POLLRDNORM);
+	insint(d, "POLLRDBAND", POLLRDBAND);
+	insint(d, "POLLWRNORM", POLLWRNORM);
+	insint(d, "POLLWRBAND", POLLWRBAND);
+	insint(d, "POLLMSG", POLLMSG);
+#endif /* HAVE_POLL */
 }