[Python-checkins] r68486 - in python/branches/py3k: Lib/gzip.py Lib/test/test_bz2.py Lib/test/test_gzip.py Misc/NEWS Modules/bz2module.c

[antoine.pitrou]

Author: antoine.pitrou
Date: Sat Jan 10 17:22:51 2009
New Revision: 68486

Log:
Merged revisions 68484-68485 via svnmerge from 
svn+ssh://pythondev@svn.python.org/python/trunk

........
  r68484 | antoine.pitrou | 2009-01-10 17:13:45 +0100 (sam., 10 janv. 2009) | 3 lines
  
  Issue #3860: GzipFile and BZ2File now support the context manager protocol.
........
  r68485 | antoine.pitrou | 2009-01-10 17:15:24 +0100 (sam., 10 janv. 2009) | 1 line
  
  Add NEWS entry for r68484.
........


Modified:
   python/branches/py3k/   (props changed)
   python/branches/py3k/Lib/gzip.py
   python/branches/py3k/Lib/test/test_bz2.py
   python/branches/py3k/Lib/test/test_gzip.py
   python/branches/py3k/Misc/NEWS
   python/branches/py3k/Modules/bz2module.c

Modified: python/branches/py3k/Lib/gzip.py
==============================================================================
--- python/branches/py3k/Lib/gzip.py	(original)
+++ python/branches/py3k/Lib/gzip.py	Sat Jan 10 17:22:51 2009
@@ -472,6 +472,14 @@
         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.

Modified: python/branches/py3k/Lib/test/test_bz2.py
==============================================================================
--- python/branches/py3k/Lib/test/test_bz2.py	(original)
+++ python/branches/py3k/Lib/test/test_bz2.py	Sat Jan 10 17:22:51 2009
@@ -258,6 +258,28 @@
         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):

Modified: python/branches/py3k/Lib/test/test_gzip.py
==============================================================================
--- python/branches/py3k/Lib/test/test_gzip.py	(original)
+++ python/branches/py3k/Lib/test/test_gzip.py	Sat Jan 10 17:22:51 2009
@@ -166,7 +166,6 @@
         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 @@
 
         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)
 

Modified: python/branches/py3k/Misc/NEWS
==============================================================================
--- python/branches/py3k/Misc/NEWS	(original)
+++ python/branches/py3k/Misc/NEWS	Sat Jan 10 17:22:51 2009
@@ -111,6 +111,8 @@
 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.
 

Modified: python/branches/py3k/Modules/bz2module.c
==============================================================================
--- python/branches/py3k/Modules/bz2module.c	(original)
+++ python/branches/py3k/Modules/bz2module.c	Sat Jan 10 17:22:51 2009
@@ -1086,6 +1086,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[] = {
@@ -1097,6 +1127,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 */
 };