#467924, patch by Alan McIntyre: Add ZipFile.extract and ZipFile.extractall.

4 files changed, 144 insertions, 1 deletions

diff --git a/Doc/library/zipfile.rst b/Doc/library/zipfile.rst
index e904a37..f1cfc1b 100644
--- a/Doc/library/zipfile.rst
+++ b/Doc/library/zipfile.rst
@@ -180,6 +180,27 @@ ZipFile Objects
    .. versionadded:: 2.6
 
 
+.. method:: ZipFile.extract(member[, path[, pwd]])
+
+   Extract a member from the archive to the current working directory, using its
+   full name.  Its file information is extracted as accurately as possible.
+   *path* specifies a different directory to extract to.   *member* can be a
+   filename or a :class:`ZipInfo` object.  *pwd* is the password used for
+   encrypted files.
+
+   .. versionadded:: 2.6
+
+
+.. method:: ZipFile.extractall([path[, members[, pwd]]])
+
+   Extract all members from the archive to the current working directory.  *path* 
+   specifies a different directory to extract to.  *members* is optional and must
+   be a subset of the list returned by :meth:`namelist`.  *pwd* is the password
+   used for encrypted files.
+
+   .. versionadded:: 2.6
+
+
 .. method:: ZipFile.printdir()
 
    Print a table of contents for the archive to ``sys.stdout``.
@@ -249,6 +270,13 @@ ZipFile Objects
    created with mode ``'r'``  will raise a :exc:`RuntimeError`.  Calling
    :meth:`writestr` on a closed ZipFile will raise a :exc:`RuntimeError`.
 
+   .. note::
+
+      When passing a :class:`ZipInfo` instance as the *zinfo_or_acrname* parameter, 
+      the compression method used will be that specified in the *compress_type* 
+      member of the given :class:`ZipInfo` instance.  By default, the 
+      :class:`ZipInfo` constructor sets this member to :const:`ZIP_STORED`.
+
 The following data attribute is also available:
 
 
diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py
index 3d2f9bd..40003aa 100644
--- a/Lib/test/test_zipfile.py
+++ b/Lib/test/test_zipfile.py
@@ -16,6 +16,11 @@ from test.test_support import TESTFN, run_unittest
 TESTFN2 = TESTFN + "2"
 FIXEDTEST_SIZE = 1000
 
+SMALL_TEST_DATA = [('_ziptest1', '1q2w3e4r5t'),
+                   ('ziptest2dir/_ziptest2', 'qawsedrftg'),
+                   ('/ziptest2dir/ziptest3dir/_ziptest3', 'azsxdcfvgb'),
+                   ('ziptest2dir/ziptest3dir/ziptest4dir/_ziptest3', '6y7u8i9o0p')]
+
 class TestsWithSourceFile(unittest.TestCase):
     def setUp(self):
         self.line_gen = ["Zipfile test line %d. random float: %f" % (i, random())
@@ -296,6 +301,57 @@ class TestsWithSourceFile(unittest.TestCase):
         self.assertRaises(RuntimeError, zipf.write, TESTFN)
         zipf.close()
 
+    def testExtract(self):
+        zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED)
+        for fpath, fdata in SMALL_TEST_DATA:
+            zipfp.writestr(fpath, fdata)
+        zipfp.close()
+
+        zipfp = zipfile.ZipFile(TESTFN2, "r")
+        for fpath, fdata in SMALL_TEST_DATA:
+            writtenfile = zipfp.extract(fpath)
+
+            # make sure it was written to the right place
+            if os.path.isabs(fpath):
+                correctfile = os.path.join(os.getcwd(), fpath[1:])
+            else:
+                correctfile = os.path.join(os.getcwd(), fpath)
+
+            self.assertEqual(writtenfile, correctfile)
+
+            # make sure correct data is in correct file
+            self.assertEqual(fdata, file(writtenfile, "rb").read())
+
+            os.remove(writtenfile)
+
+        zipfp.close()
+
+        # remove the test file subdirectories
+        shutil.rmtree(os.path.join(os.getcwd(), 'ziptest2dir'))
+
+    def testExtractAll(self):
+        zipfp = zipfile.ZipFile(TESTFN2, "w", zipfile.ZIP_STORED)
+        for fpath, fdata in SMALL_TEST_DATA:
+            zipfp.writestr(fpath, fdata)
+        zipfp.close()
+
+        zipfp = zipfile.ZipFile(TESTFN2, "r")
+        zipfp.extractall()
+        for fpath, fdata in SMALL_TEST_DATA:
+            if os.path.isabs(fpath):
+                outfile = os.path.join(os.getcwd(), fpath[1:])
+            else:
+                outfile = os.path.join(os.getcwd(), fpath)
+
+            self.assertEqual(fdata, file(outfile, "rb").read())
+
+            os.remove(outfile)
+
+        zipfp.close()
+
+        # remove the test file subdirectories
+        shutil.rmtree(os.path.join(os.getcwd(), 'ziptest2dir'))
+
     def tearDown(self):
         os.remove(TESTFN)
         os.remove(TESTFN2)
diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index a53f8ed..ab9c93f 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -1,7 +1,7 @@
 """
 Read and write ZIP files.
 """
-import struct, os, time, sys
+import struct, os, time, sys, shutil
 import binascii, cStringIO
 
 try:
@@ -807,6 +807,62 @@ class ZipFile:
             zef.set_univ_newlines(True)
         return zef
 
+    def extract(self, member, path=None, pwd=None):
+        """Extract a member from the archive to the current working directory,
+           using its full name. Its file information is extracted as accurately
+           as possible. `member' may be a filename or a ZipInfo object. You can
+           specify a different directory using `path'.
+        """
+        if not isinstance(member, ZipInfo):
+            member = self.getinfo(member)
+
+        if path is None:
+            path = os.getcwd()
+
+        return self._extract_member(member, path, pwd)
+
+    def extractall(self, path=None, members=None, pwd=None):
+        """Extract all members from the archive to the current working
+           directory. `path' specifies a different directory to extract to.
+           `members' is optional and must be a subset of the list returned
+           by namelist().
+        """
+        if members is None:
+            members = self.namelist()
+
+        for zipinfo in members:
+            self.extract(zipinfo, path, pwd)
+
+    def _extract_member(self, member, targetpath, pwd):
+        """Extract the ZipInfo object 'member' to a physical
+           file on the path targetpath.
+        """
+        # build the destination pathname, replacing
+        # forward slashes to platform specific separators.
+        if targetpath[-1:] == "/":
+            targetpath = targetpath[:-1]
+
+        # don't include leading "/" from file name if present
+        if os.path.isabs(member.filename):
+            targetpath = os.path.join(targetpath, member.filename[1:])
+        else:
+            targetpath = os.path.join(targetpath, member.filename)
+
+        targetpath = os.path.normpath(targetpath)
+
+        # Create all upper directories if necessary.
+        upperdirs = os.path.dirname(targetpath)
+        if upperdirs and not os.path.exists(upperdirs):
+            os.makedirs(upperdirs)
+
+        source = self.open(member.filename, pwd=pwd)
+        target = file(targetpath, "wb")
+        shutil.copyfileobj(source, target)
+        source.close()
+        target.close()
+
+        return targetpath
+
     def _writecheck(self, zinfo):
         """Check for errors before writing a file to the archive."""
         if zinfo.filename in self.NameToInfo:
diff --git a/Misc/NEWS b/Misc/NEWS
index e76e733..4699e1b 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -348,6 +348,9 @@ Core and builtins
 Library
 -------
 
+- Patch #467924: add ZipFile.extract() and ZipFile.extractall() in the
+  zipfile module.
+
 - Issue #1646: Make socket support the TIPC protocol.
 
 - Bug #1742: return os.curdir from os.path.relpath() if both arguments are