#8046: add context manager protocol support to mmap objects. Also add closed property.

1 files changed, 34 insertions, 1 deletions

diff --git a/Modules/mmapmodule.c b/Modules/mmapmodule.c
index 9714ddf..de527ce 100644
--- a/Modules/mmapmodule.c
+++ b/Modules/mmapmodule.c
@@ -651,6 +651,31 @@ mmap_move_method(mmap_object *self, PyObject *args)
     }
 }
 
+static PyObject *
+mmap_closed_get(mmap_object *self)
+{
+#ifdef MS_WINDOWS
+    return PyBool_FromLong(self->map_handle == NULL ? 1 : 0);
+#elif defined(UNIX)
+    return PyBool_FromLong(self->data == NULL ? 1 : 0);
+#endif
+}
+
+static PyObject *
+mmap__enter__method(mmap_object *self, PyObject *args)
+{
+    CHECK_VALID(NULL);
+
+    Py_INCREF(self);
+    return (PyObject *)self;
+}
+
+static PyObject *
+mmap__exit__method(PyObject *self, PyObject *args)
+{
+    return PyObject_CallMethod(self, "close", NULL);
+}
+
 static struct PyMethodDef mmap_object_methods[] = {
     {"close",           (PyCFunction) mmap_close_method,        METH_NOARGS},
     {"find",            (PyCFunction) mmap_find_method,         METH_VARARGS},
@@ -666,9 +691,17 @@ static struct PyMethodDef mmap_object_methods[] = {
     {"tell",            (PyCFunction) mmap_tell_method,         METH_NOARGS},
     {"write",           (PyCFunction) mmap_write_method,        METH_VARARGS},
     {"write_byte",      (PyCFunction) mmap_write_byte_method,   METH_VARARGS},
+    {"__enter__",       (PyCFunction) mmap__enter__method,      METH_NOARGS},
+    {"__exit__",        (PyCFunction) mmap__exit__method,       METH_VARARGS},
     {NULL,         NULL}       /* sentinel */
 };
 
+static PyGetSetDef mmap_object_getset[] = {
+    {"closed", (getter) mmap_closed_get, NULL, NULL},
+    {NULL}
+};
+
+
 /* Functions for treating an mmap'ed file as a buffer */
 
 static int
@@ -975,7 +1008,7 @@ static PyTypeObject mmap_object_type = {
     0,                                          /* tp_iternext */
     mmap_object_methods,                        /* tp_methods */
     0,                                          /* tp_members */
-    0,                                          /* tp_getset */
+    mmap_object_getset,                         /* tp_getset */
     0,                                          /* tp_base */
     0,                                          /* tp_dict */
     0,                                          /* tp_descr_get */