Issue #16488: epoll() objects now support the with statement.

pitrou · pitrou · commit 09bb89b8cf24 · 2012-12-15T21:14:21.000+01:00
Patch by Serhiy Storchaka.
diff --git a/Doc/library/select.rst b/Doc/library/select.rst
@@ -47,11 +47,14 @@ The module defines the following:
    to :const:`EPOLL_CLOEXEC`, which causes the epoll descriptor to be closed
    automatically when :func:`os.execve` is called. See section
    :ref:`epoll-objects` below for the methods supported by epolling objects.
-
+   They also support the :keyword:`with` statement.
 
    .. versionchanged:: 3.3
       Added the *flags* parameter.
 
+   .. versionchanged:: 3.4
+      Support for the :keyword:`with` statement was added.
+
 
 .. function:: poll()
 
diff --git a/Lib/test/test_epoll.py b/Lib/test/test_epoll.py
@@ -87,6 +87,13 @@ def test_badcreate(self):
         self.assertRaises(TypeError, select.epoll, ['foo'])
         self.assertRaises(TypeError, select.epoll, {})
 
+    def test_context_manager(self):
+        with select.epoll(16) as ep:
+            self.assertGreater(ep.fileno(), 0)
+            self.assertFalse(ep.closed)
+        self.assertTrue(ep.closed)
+        self.assertRaises(ValueError, ep.fileno)
+
     def test_add(self):
         server, client = self._connected_pair()
 
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -167,6 +167,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #16488: epoll() objects now support the `with` statement.  Patch
+  by Serhiy Storchaka.
+
 - Issue #16298: In HTTPResponse.read(), close the socket when there is no
   Content-Length and the incoming stream is finished.  Patch by Eran
   Rundstein.
diff --git a/Modules/selectmodule.c b/Modules/selectmodule.c
@@ -1394,6 +1394,24 @@ Wait for events on the epoll file descriptor for a maximum time of timeout\n\
 in seconds (as float). -1 makes poll wait indefinitely.\n\
 Up to maxevents are returned to the caller.");
 
+static PyObject *
+pyepoll_enter(pyEpoll_Object *self, PyObject *args)
+{
+    if (self->epfd < 0)
+        return pyepoll_err_closed();
+
+    Py_INCREF(self);
+    return (PyObject *)self;
+}
+
+static PyObject *
+pyepoll_exit(PyObject *self, PyObject *args)
+{
+    _Py_IDENTIFIER(close);
+
+    return _PyObject_CallMethodId(self, &PyId_close, NULL);
+}
+
 static PyMethodDef pyepoll_methods[] = {
     {"fromfd",          (PyCFunction)pyepoll_fromfd,
      METH_VARARGS | METH_CLASS, pyepoll_fromfd_doc},
@@ -1409,6 +1427,10 @@ static PyMethodDef pyepoll_methods[] = {
      METH_VARARGS | METH_KEYWORDS,      pyepoll_unregister_doc},
     {"poll",            (PyCFunction)pyepoll_poll,
      METH_VARARGS | METH_KEYWORDS,      pyepoll_poll_doc},
+    {"__enter__",           (PyCFunction)pyepoll_enter,     METH_NOARGS,
+     NULL},
+    {"__exit__",           (PyCFunction)pyepoll_exit,     METH_VARARGS,
+     NULL},
     {NULL,      NULL},
 };
 
