Merge #14399: zipfile now correctly handles comments added to empty zipfiles.

Patch by Serhiy Storchaka.

This also moves the TypeError that results from trying to use a unicode
comment from the 'close' step to the point at which the comment is added to
the zipfile.

1 files changed, 22 insertions, 11 deletions

diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index 6ca269f..435a3e9 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -698,7 +698,7 @@ class ZipFile:
         self.compression = compression  # Method of compression
         self.mode = key = mode.replace('b', '')[0]
         self.pwd = None
-        self.comment = b''
+        self._comment = b''
 
         # Check if we were passed a file-like object
         if isinstance(file, str):
@@ -774,7 +774,7 @@ class ZipFile:
             print(endrec)
         size_cd = endrec[_ECD_SIZE]             # bytes in central directory
         offset_cd = endrec[_ECD_OFFSET]         # offset of central directory
-        self.comment = endrec[_ECD_COMMENT]     # archive comment
+        self._comment = endrec[_ECD_COMMENT]    # archive comment
 
         # "concat" is zero, unless zip was concatenated to another file
         concat = endrec[_ECD_LOCATION] - size_cd - offset_cd
@@ -886,6 +886,24 @@ class ZipFile:
         else:
             self.pwd = None
 
+    @property
+    def comment(self):
+        """The comment text associated with the ZIP file."""
+        return self._comment
+
+    @comment.setter
+    def comment(self, comment):
+        if not isinstance(comment, bytes):
+            raise TypeError("comment: expected bytes, got %s" % type(comment))
+        # check for valid comment length
+        if len(comment) >= ZIP_MAX_COMMENT:
+            if self.debug:
+                print('Archive comment is too long; truncating to %d bytes'
+                        % ZIP_MAX_COMMENT)
+            comment = comment[:ZIP_MAX_COMMENT]
+        self._comment = comment
+        self._didModify = True
+
     def read(self, name, pwd=None):
         """Return file bytes (as a string) for name."""
         with self.open(name, "r", pwd) as fp:
@@ -1287,18 +1305,11 @@ class ZipFile:
                 centDirSize = min(centDirSize, 0xFFFFFFFF)
                 centDirOffset = min(centDirOffset, 0xFFFFFFFF)
 
-            # check for valid comment length
-            if len(self.comment) >= ZIP_MAX_COMMENT:
-                if self.debug > 0:
-                    msg = 'Archive comment is too long; truncating to %d bytes' \
-                          % ZIP_MAX_COMMENT
-                self.comment = self.comment[:ZIP_MAX_COMMENT]
-
             endrec = struct.pack(structEndArchive, stringEndArchive,
                                  0, 0, centDirCount, centDirCount,
-                                 centDirSize, centDirOffset, len(self.comment))
+                                 centDirSize, centDirOffset, len(self._comment))
             self.fp.write(endrec)
-            self.fp.write(self.comment)
+            self.fp.write(self._comment)
             self.fp.flush()
 
         if not self._filePassed: