diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/test/test_zipfile.py | 37 | ||||
| -rw-r--r-- | Lib/zipfile.py | 39 |
2 files changed, 58 insertions, 18 deletions
diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py index 0dc080b..8da74f5 100644 --- a/Lib/test/test_zipfile.py +++ b/Lib/test/test_zipfile.py @@ -1,24 +1,41 @@ -import zipfile, os +import zipfile, os, StringIO, tempfile from test_support import TestFailed srcname = "junk9630.tmp" zipname = "junk9708.tmp" + +def zipTest(f, compression, srccontents): + zip = zipfile.ZipFile(f, "w", compression) # Create the ZIP archive + zip.write(srcname, "another.name") + zip.write(srcname, srcname) + zip.close() + + zip = zipfile.ZipFile(f, "r", compression) # Read the ZIP archive + readData2 = zip.read(srcname) + readData1 = zip.read("another.name") + zip.close() + + if readData1 != srccontents or readData2 != srccontents: + raise TestFailed, "Written data doesn't equal read data." + + try: - fp = open(srcname, "w") # Make a source file with some lines + fp = open(srcname, "wb") # Make a source file with some lines for i in range(0, 1000): fp.write("Test of zipfile line %d.\n" % i) fp.close() + + fp = open(srcname, "rb") + writtenData = fp.read() + fp.close() + + for file in (zipname, tempfile.TemporaryFile(), StringIO.StringIO()): + zipTest(file, zipfile.ZIP_STORED, writtenData) - zip = zipfile.ZipFile(zipname, "w") # Create the ZIP archive - zip.write(srcname, srcname) - zip.write(srcname, "another.name") - zip.close() + for file in (zipname, tempfile.TemporaryFile(), StringIO.StringIO()): + zipTest(file, zipfile.ZIP_DEFLATED, writtenData) - zip = zipfile.ZipFile(zipname, "r") # Read the ZIP archive - zip.read("another.name") - zip.read(srcname) - zip.close() finally: if os.path.isfile(srcname): # Remove temporary files os.unlink(srcname) diff --git a/Lib/zipfile.py b/Lib/zipfile.py index 39b1511..bf59043 100644 --- a/Lib/zipfile.py +++ b/Lib/zipfile.py @@ -65,6 +65,9 @@ _FH_UNCOMPRESSED_SIZE = 9 _FH_FILENAME_LENGTH = 10 _FH_EXTRA_FIELD_LENGTH = 11 +# Used to compare file passed to ZipFile +_STRING_TYPES = (type('s'), type(u's')) + def is_zipfile(filename): """Quickly see if file is a ZIP file by checking the magic number. @@ -128,11 +131,19 @@ class ZipInfo: class ZipFile: - """Class with methods to open, read, write, close, list zip files.""" + """ Class with methods to open, read, write, close, list zip files. + + z = ZipFile(file, mode="r", compression=ZIP_STORED) + + file: Either the path to the file, or a file-like object. + If it is a path, the file will be opened and closed by ZipFile. + mode: The mode can be either read "r", write "w" or append "a". + compression: ZIP_STORED (no compression) or ZIP_DEFLATED (requires zlib). + """ fp = None # Set here since __del__ checks it - def __init__(self, filename, mode="r", compression=ZIP_STORED): + def __init__(self, file, mode="r", compression=ZIP_STORED): """Open the ZIP file with mode read "r", write "w" or append "a".""" if compression == ZIP_STORED: pass @@ -146,15 +157,25 @@ class ZipFile: self.NameToInfo = {} # Find file info given name self.filelist = [] # List of ZipInfo instances for archive self.compression = compression # Method of compression - self.filename = filename self.mode = key = mode[0] + + # Check if we were passed a file-like object + if type(file) in _STRING_TYPES: + self._filePassed = 0 + self.filename = file + modeDict = {'r' : 'rb', 'w': 'wb', 'a' : 'r+b'} + self.fp = open(file, modeDict[mode]) + else: + self._filePassed = 1 + self.fp = file + self.filename = getattr(file, 'name', None) + if key == 'r': - self.fp = open(filename, "rb") self._GetContents() elif key == 'w': - self.fp = open(filename, "wb") + pass elif key == 'a': - fp = self.fp = open(filename, "r+b") + fp = self.fp fp.seek(-22, 2) # Seek to end-of-file record endrec = fp.read() if endrec[0:4] == stringEndArchive and \ @@ -401,7 +422,7 @@ class ZipFile: def __del__(self): """Call the "close()" method in case the user forgot.""" - if self.fp: + if self.fp and not self._filePassed: self.fp.close() self.fp = None @@ -433,7 +454,9 @@ class ZipFile: endrec = struct.pack(structEndArchive, stringEndArchive, 0, 0, count, count, pos2 - pos1, pos1, 0) self.fp.write(endrec) - self.fp.close() + self.fp.flush() + if not self._filePassed: + self.fp.close() self.fp = None |