Issue 6003: ZipFile.writestr "compression_type" argument

diff --git a/Doc/library/zipfile.rst b/Doc/library/zipfile.rst
index 0d5f2054047cccd7943d7edce6f162da1b11afcf..cbaf0a6a8133f24edb79ec68f31c5aa90c13d887 100644 (file)
--- a/Doc/library/zipfile.rst
+++ b/Doc/library/zipfile.rst
@@ -284,7 +284,7 @@ ZipFile Objects
       byte, the name of the file in the archive will be truncated at the null byte.
 
 
-.. method:: ZipFile.writestr(zinfo_or_arcname, bytes)
+.. method:: ZipFile.writestr(zinfo_or_arcname, bytes[, compress_type])
 
    Write the string *bytes* to the archive; *zinfo_or_arcname* is either the file
    name it will be given in the archive, or a :class:`ZipInfo` instance.  If it's
@@ -294,6 +294,10 @@ ZipFile Objects
    created with mode ``'r'``  will raise a :exc:`RuntimeError`.  Calling
    :meth:`writestr` on a closed ZipFile will raise a :exc:`RuntimeError`.
 
+   If given, *compress_type* overrides the value given for the *compression*
+   parameter to the constructor for the new entry, or in the *zinfo_or_arcname*
+   (if that is a :class:`ZipInfo` instance).
+
    .. note::
 
       When passing a :class:`ZipInfo` instance as the *zinfo_or_acrname* parameter,
@@ -301,6 +305,9 @@ ZipFile Objects
       member of the given :class:`ZipInfo` instance.  By default, the
       :class:`ZipInfo` constructor sets this member to :const:`ZIP_STORED`.
 
+   .. versionchanged:: 2.7
+      The *compression_type* argument.
+
 The following data attributes are also available:
 
 

diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py
index 3c3bae2c2b1df746a20236cfd27bfe1e47ecbaec..16acabcf057fd82d0e6ed9fec1daf4ab0d3d907e 100644 (file)
--- a/Lib/test/test_zipfile.py
+++ b/Lib/test/test_zipfile.py
@@ -408,6 +408,20 @@ class TestsWithSourceFile(unittest.TestCase):
         # 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 3a44d3a9de43b7f46520e682c8dae255ce9ef8c6..ece538edc2c07daf263340585159eae4d2b14492 100644 (file)
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -1063,13 +1063,14 @@ class ZipFile:
         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 @@ class ZipFile:
             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)

diff --git a/Misc/NEWS b/Misc/NEWS
index 68cd072ab83106691dcd393865fc70df460d62aa..1a752fc7d8cfb4baaec512632cdc4ffcb2df2cca 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -15,6 +15,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #6003: add an argument to ``zipfile.Zipfile.writestr`` to
+  specify the compression type.
+
 
 What's New in Python 2.7 alpha 3?
 =================================