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

diff --git a/Lib/gzip.py b/Lib/gzip.py
index 931c9eff20507ace5c8dcf982db21425c0aaac0c..f568796f7a786369f178654b0e1a273437d45b50 100644 (file)
--- a/Lib/gzip.py
+++ b/Lib/gzip.py
@@ -454,6 +454,14 @@ class GzipFile:
         else:
             raise StopIteration
 
+    def __enter__(self):
+        if self.fileobj is None:
+            raise ValueError("I/O operation on closed GzipFile object")
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
 
 def _test():
     # Act like gzip; with -d, act like gunzip.

diff --git a/Lib/test/test_bz2.py b/Lib/test/test_bz2.py
index c4d9b6968a3426c21e06cccb311dc88f2e054bbf..2f627c32011d9736df9fe6251d85d98147e9c449 100644 (file)
--- a/Lib/test/test_bz2.py
+++ b/Lib/test/test_bz2.py
@@ -284,6 +284,28 @@ class BZ2FileTest(BaseTest):
         bz2f.close()
         self.assertEqual(xlines, ['Test'])
 
+    def testContextProtocol(self):
+        # BZ2File supports the context management protocol
+        f = None
+        with BZ2File(self.filename, "wb") as f:
+            f.write(b"xxx")
+        f = BZ2File(self.filename, "rb")
+        f.close()
+        try:
+            with f:
+                pass
+        except ValueError:
+            pass
+        else:
+            self.fail("__enter__ on a closed file didn't raise an exception")
+        try:
+            with BZ2File(self.filename, "wb") as f:
+                1/0
+        except ZeroDivisionError:
+            pass
+        else:
+            self.fail("1/0 didn't raise an exception")
+
 
 class BZ2CompressorTest(BaseTest):
     def testCompress(self):

diff --git a/Lib/test/test_gzip.py b/Lib/test/test_gzip.py
index 25f9ee628503cdb3319ee24321475c1004b3ee10..b7250dbddfd788a70fb68fe01e3f0c1cc4c75ae5 100644 (file)
--- a/Lib/test/test_gzip.py
+++ b/Lib/test/test_gzip.py
@@ -166,7 +166,6 @@ class TestGzip(unittest.TestCase):
         fWrite = gzip.GzipFile(self.filename, 'w', mtime = mtime)
         fWrite.write(data1)
         fWrite.close()
-
         fRead = gzip.GzipFile(self.filename)
         dataRead = fRead.read()
         self.assertEqual(dataRead, data1)
@@ -222,6 +221,27 @@ class TestGzip(unittest.TestCase):
 
         fRead.close()
 
+    def test_with_open(self):
+        # GzipFile supports the context management protocol
+        with gzip.GzipFile(self.filename, "wb") as f:
+            f.write(b"xxx")
+        f = gzip.GzipFile(self.filename, "rb")
+        f.close()
+        try:
+            with f:
+                pass
+        except ValueError:
+            pass
+        else:
+            self.fail("__enter__ on a closed file didn't raise an exception")
+        try:
+            with gzip.GzipFile(self.filename, "wb") as f:
+                1/0
+        except ZeroDivisionError:
+            pass
+        else:
+            self.fail("1/0 didn't raise an exception")
+
 def test_main(verbose=None):
     test_support.run_unittest(TestGzip)
 

diff --git a/Modules/bz2module.c b/Modules/bz2module.c
index b5542bf968d3430a241707586f92898f7f2de10a..045d7b270c03d890741971465dcca907735ff468 100644 (file)
--- 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 */
 };