Itamar Shtull-Trauring <itamar@maxnm.com>:
Add support to zipfile to support opening an archive represented by an
open file rather than a file name.
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_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 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 @@
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 @@
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 @@
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