From 0138581c43c9812d9bd415d95523e65fd72e6a34 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lars=20Gust=C3=A4bel?= <lars@gustaebel.de>
Date: Wed, 3 Mar 2010 12:08:54 +0000
Subject: [PATCH] Merged revisions 78623 via svnmerge from
 svn+ssh://pythondev@svn.python.org/python/trunk

........
  r78623 | lars.gustaebel | 2010-03-03 12:55:48 +0100 (Wed, 03 Mar 2010) | 3 lines

  Issue #7232: Add support for the context manager protocol
  to the TarFile class.
........
---
 Doc/library/tarfile.rst  | 15 ++++++++++
 Lib/tarfile.py           | 14 ++++++++++
 Lib/test/test_tarfile.py | 60 ++++++++++++++++++++++++++++++++++++++++
 Misc/NEWS                |  3 ++
 4 files changed, 92 insertions(+)

diff --git a/Doc/library/tarfile.rst b/Doc/library/tarfile.rst
index 4d0a99527aa..8b53b578247 100644
--- a/Doc/library/tarfile.rst
+++ b/Doc/library/tarfile.rst
@@ -209,6 +209,14 @@ a header block followed by data blocks. It is possible to store a file in a tar
 archive several times. Each archive member is represented by a :class:`TarInfo`
 object, see :ref:`tarinfo-objects` for details.
 
+A :class:`TarFile` object can be used as a context manager in a :keyword:`with`
+statement. It will automatically be closed when the block is completed. Please
+note that in the event of an exception an archive opened for writing will not
+be finalized, only the internally used file object will be closed. See the
+:ref:`tar-examples` section for a use case.
+
+.. versionadded:: 3.2
+   Added support for the context manager protocol.
 
 .. class:: TarFile(name=None, mode='r', fileobj=None, format=DEFAULT_FORMAT, tarinfo=TarInfo, dereference=False, ignore_zeros=False, encoding=ENCODING, errors=None, pax_headers=None, debug=0, errorlevel=0)
 
@@ -593,6 +601,13 @@ How to create an uncompressed tar archive from a list of filenames::
        tar.add(name)
    tar.close()
 
+The same example using the :keyword:`with` statement::
+
+    import tarfile
+    with tarfile.open("sample.tar", "w") as tar:
+        for name in ["foo", "bar", "quux"]:
+            tar.add(name)
+
 How to read a gzip compressed tar archive and display some member information::
 
    import tarfile
diff --git a/Lib/tarfile.py b/Lib/tarfile.py
index 7f08a984061..e28b1c7e578 100644
--- a/Lib/tarfile.py
+++ b/Lib/tarfile.py
@@ -2391,6 +2391,20 @@ class TarFile(object):
         """
         if level <= self.debug:
             print(msg, file=sys.stderr)
+
+    def __enter__(self):
+        self._check()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        if type is None:
+            self.close()
+        else:
+            # An exception occurred. We must not call close() because
+            # it would try to write end-of-archive blocks and padding.
+            if not self._extfileobj:
+                self.fileobj.close()
+            self.closed = True
 # class TarFile
 
 class TarIter:
diff --git a/Lib/test/test_tarfile.py b/Lib/test/test_tarfile.py
index 792f2b93a91..12cf2fbffdb 100644
--- a/Lib/test/test_tarfile.py
+++ b/Lib/test/test_tarfile.py
@@ -1274,6 +1274,65 @@ class MiscTest(unittest.TestCase):
         self.assertEqual(tarfile.itn(0xffffffff), b"\x80\x00\x00\x00\xff\xff\xff\xff")
 
 
+class ContextManagerTest(unittest.TestCase):
+
+    def test_basic(self):
+        with tarfile.open(tarname) as tar:
+            self.assertFalse(tar.closed, "closed inside runtime context")
+        self.assertTrue(tar.closed, "context manager failed")
+
+    def test_closed(self):
+        # The __enter__() method is supposed to raise IOError
+        # if the TarFile object is already closed.
+        tar = tarfile.open(tarname)
+        tar.close()
+        with self.assertRaises(IOError):
+            with tar:
+                pass
+
+    def test_exception(self):
+        # Test if the IOError exception is passed through properly.
+        with self.assertRaises(Exception) as exc:
+            with tarfile.open(tarname) as tar:
+                raise IOError
+        self.assertIsInstance(exc.exception, IOError,
+                              "wrong exception raised in context manager")
+        self.assertTrue(tar.closed, "context manager failed")
+
+    def test_no_eof(self):
+        # __exit__() must not write end-of-archive blocks if an
+        # exception was raised.
+        try:
+            with tarfile.open(tmpname, "w") as tar:
+                raise Exception
+        except:
+            pass
+        self.assertEqual(os.path.getsize(tmpname), 0,
+                "context manager wrote an end-of-archive block")
+        self.assertTrue(tar.closed, "context manager failed")
+
+    def test_eof(self):
+        # __exit__() must write end-of-archive blocks, i.e. call
+        # TarFile.close() if there was no error.
+        with tarfile.open(tmpname, "w"):
+            pass
+        self.assertNotEqual(os.path.getsize(tmpname), 0,
+                "context manager wrote no end-of-archive block")
+
+    def test_fileobj(self):
+        # Test that __exit__() did not close the external file
+        # object.
+        fobj = open(tmpname, "wb")
+        try:
+            with tarfile.open(fileobj=fobj, mode="w") as tar:
+                raise Exception
+        except:
+            pass
+        self.assertFalse(fobj.closed, "external file object was closed")
+        self.assertTrue(tar.closed, "context manager failed")
+        fobj.close()
+
+
 class GzipMiscReadTest(MiscReadTest):
     tarname = gzipname
     mode = "r:gz"
@@ -1354,6 +1413,7 @@ def test_main():
         AppendTest,
         LimitsTest,
         MiscTest,
+        ContextManagerTest,
     ]
 
     if hasattr(os, "link"):
diff --git a/Misc/NEWS b/Misc/NEWS
index e8f37ab8b96..9fd6a9a8b14 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -265,6 +265,9 @@ C-API
 Library
 -------
 
+- Issue #7232: Add support for the context manager protocol to the TarFile
+  class.
+
 - Issue #7250: Fix info leak of os.environ across multi-run uses of
   wsgiref.handlers.CGIHandler.