diff options
author | Georg Brandl <georg@python.org> | 2008-01-07 18:47:44 (GMT) |
---|---|---|
committer | Georg Brandl <georg@python.org> | 2008-01-07 18:47:44 (GMT) |
commit | 62416bcf5ad46019f758aa212b05b3d60cb71ef3 (patch) | |
tree | 168c2f432e63ebc9265a93e41274c5ca53513a48 /Lib | |
parent | 76b30d1688a7ba1ff1b01a3eb21bf4890f71d404 (diff) | |
download | cpython-62416bcf5ad46019f758aa212b05b3d60cb71ef3.zip cpython-62416bcf5ad46019f758aa212b05b3d60cb71ef3.tar.gz cpython-62416bcf5ad46019f758aa212b05b3d60cb71ef3.tar.bz2 |
#467924, patch by Alan McIntyre: Add ZipFile.extract and ZipFile.extractall.
Diffstat (limited to 'Lib')
-rw-r--r-- | Lib/test/test_zipfile.py | 56 | ||||
-rw-r--r-- | Lib/zipfile.py | 58 |
2 files changed, 113 insertions, 1 deletions
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: |