Issue #3860: GzipFile and BZ2File now support the context manager protocol.
diff --git a/Modules/bz2module.c b/Modules/bz2module.c
index b5542bf..045d7b2 100644
--- a/Modules/bz2module.c
+++ b/Modules/bz2module.c
@@ -1201,6 +1201,36 @@
return ret;
}
+PyDoc_STRVAR(BZ2File_enter_doc,
+"__enter__() -> self.");
+
+static PyObject *
+BZ2File_enter(BZ2FileObject *self)
+{
+ if (self->mode == MODE_CLOSED) {
+ PyErr_SetString(PyExc_ValueError,
+ "I/O operation on closed file");
+ return NULL;
+ }
+ Py_INCREF(self);
+ return (PyObject *) self;
+}
+
+PyDoc_STRVAR(BZ2File_exit_doc,
+"__exit__(*excinfo) -> None. Closes the file.");
+
+static PyObject *
+BZ2File_exit(BZ2FileObject *self, PyObject *args)
+{
+ PyObject *ret = PyObject_CallMethod((PyObject *) self, "close", NULL);
+ if (!ret)
+ /* If error occurred, pass through */
+ return NULL;
+ Py_DECREF(ret);
+ Py_RETURN_NONE;
+}
+
+
static PyObject *BZ2File_getiter(BZ2FileObject *self);
static PyMethodDef BZ2File_methods[] = {
@@ -1213,6 +1243,8 @@
{"seek", (PyCFunction)BZ2File_seek, METH_VARARGS, BZ2File_seek__doc__},
{"tell", (PyCFunction)BZ2File_tell, METH_NOARGS, BZ2File_tell__doc__},
{"close", (PyCFunction)BZ2File_close, METH_NOARGS, BZ2File_close__doc__},
+ {"__enter__", (PyCFunction)BZ2File_enter, METH_NOARGS, BZ2File_enter_doc},
+ {"__exit__", (PyCFunction)BZ2File_exit, METH_VARARGS, BZ2File_exit_doc},
{NULL, NULL} /* sentinel */
};