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

1 files changed, 22 insertions, 0 deletions

diff --git a/Lib/test/test_bz2.py b/Lib/test/test_bz2.py
index c4d9b69..2f627c3 100644
--- 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):