diff options
| -rw-r--r-- | Lib/gzip.py | 8 | ||||
| -rw-r--r-- | Lib/test/test_bz2.py | 22 | ||||
| -rw-r--r-- | Lib/test/test_gzip.py | 22 | ||||
| -rw-r--r-- | Misc/NEWS | 2 | ||||
| -rw-r--r-- | Modules/bz2module.c | 32 |
5 files changed, 85 insertions, 1 deletions
diff --git a/Lib/gzip.py b/Lib/gzip.py index 560a722..019c3e2 100644 --- a/Lib/gzip.py +++ b/Lib/gzip.py @@ -472,6 +472,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 366ab7a..18d84bb 100644 --- a/Lib/test/test_bz2.py +++ b/Lib/test/test_bz2.py @@ -258,6 +258,28 @@ class BZ2FileTest(BaseTest): bz2f.close() self.assertEqual(xlines, [b'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 e758826..baf3a21 100644 --- 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) @@ -223,6 +222,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): support.run_unittest(TestGzip) @@ -111,6 +111,8 @@ Core and Builtins Library ------- +- Issue #3860: GzipFile and BZ2File now support the context manager protocol. + - Issue #4867: Fixed a crash in ctypes when passing a string to a function without defining argtypes. diff --git a/Modules/bz2module.c b/Modules/bz2module.c index 4bf31ee..c7cc767 100644 --- a/Modules/bz2module.c +++ b/Modules/bz2module.c @@ -1086,6 +1086,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[] = { @@ -1097,6 +1127,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 */ }; |