Issue #3860: GzipFile and BZ2File now support the context manager protocol.

1 files changed, 32 insertions, 0 deletions

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 @@ BZ2File_close(BZ2FileObject *self)
 	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 @@ static PyMethodDef BZ2File_methods[] = {
 	{"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 */
 };