Issue 6003: ZipFile.writestr "compression_type" argument
diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py
index 3c3bae2..16acabc 100644
--- a/Lib/test/test_zipfile.py
+++ b/Lib/test/test_zipfile.py
@@ -408,6 +408,20 @@
# remove the test file subdirectories
shutil.rmtree(os.path.join(os.getcwd(), 'ziptest2dir'))
+ def test_writestr_compression(self):
+ zipfp = zipfile.ZipFile(TESTFN2, "w")
+ zipfp.writestr("a.txt", "hello world", compress_type=zipfile.ZIP_STORED)
+ if zlib:
+ zipfp.writestr("b.txt", "hello world", compress_type=zipfile.ZIP_DEFLATED)
+
+ info = zipfp.getinfo('a.txt')
+ self.assertEqual(info.compress_type, zipfile.ZIP_STORED)
+
+ if zlib:
+ info = zipfp.getinfo('b.txt')
+ self.assertEqual(info.compress_type, zipfile.ZIP_DEFLATED)
+
+
def zip_test_writestr_permissions(self, f, compression):
# Make sure that writestr creates files with mode 0600,
# when it is passed a name rather than a ZipInfo instance.
diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index 3a44d3a..ece538e 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -1063,13 +1063,14 @@
self.filelist.append(zinfo)
self.NameToInfo[zinfo.filename] = zinfo
- def writestr(self, zinfo_or_arcname, bytes):
+ def writestr(self, zinfo_or_arcname, bytes, compress_type=None):
"""Write a file into the archive. The contents is the string
'bytes'. 'zinfo_or_arcname' is either a ZipInfo instance or
the name of the file in the archive."""
if not isinstance(zinfo_or_arcname, ZipInfo):
zinfo = ZipInfo(filename=zinfo_or_arcname,
date_time=time.localtime(time.time())[:6])
+
zinfo.compress_type = self.compression
zinfo.external_attr = 0600 << 16
else:
@@ -1079,6 +1080,9 @@
raise RuntimeError(
"Attempt to write to ZIP archive that was already closed")
+ if compress_type is not None:
+ zinfo.compress_type = compress_type
+
zinfo.file_size = len(bytes) # Uncompressed size
zinfo.header_offset = self.fp.tell() # Start of header bytes
self._writecheck(zinfo)