3 files changed, 51 insertions, 1 deletions
diff --git a/Lib/gzip.py b/Lib/gzip.py
index 560a722bae..019c3e2df1 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 366ab7a16f..18d84bb60e 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 e758826fa7..baf3a21f87 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)