#467924, patch by Alan McIntyre: Add ZipFile.extract and ZipFile.extractall.
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 @@
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 @@
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 @@
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: