Zip64 iterable tests + fix

tests/test_zipstream.py

@@ -96,5 +96,104 @@ def test_write_iterable_no_archive(self):
         z = zipstream.ZipFile(mode='w')
         self.assertRaises(TypeError, z.write_iter, iterable=range(10))
 
+    def test_write_iterable_zip64_with_not_allow_zip64_many_smalls(self):
+        # check many small streams that sum length require ZIP64 extensions when not allowed zip64
+        z = zipstream.ZipFile(mode='w', allowZip64=False)
+
+        def string_small_generator():
+            counter = 0
+            sample = b'zipstream0' * 10000000
+            len_sample = len(sample)
+            while counter + len_sample < zipstream.ZIP64_LIMIT:
+                counter += len_sample
+                yield sample
+
+        data = [string_small_generator(), string_small_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        try:
+            self.assertRaises(zipfile.LargeZipFile, lambda: [f.write(c) for c in z])
+            f.close()
+        except Exception:
+            raise
+        finally:
+            os.remove(f.name)
+
+    def test_write_iterable_zip64_with_not_allow_zip64_1_big_file(self):
+        # check 1 big stream that length require ZIP64 extensions when not allowed zip64
+        z = zipstream.ZipFile(mode='w', allowZip64=False)
+
+        def string_big_generator():
+            counter = 0
+            sample = b'zipstream0' * 10000000
+            len_sample = len(sample)
+            while counter < zipstream.ZIP64_LIMIT:
+                counter += len_sample
+                yield sample
+
+        data = [string_big_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        try:
+            self.assertRaises(zipfile.LargeZipFile, lambda: [f.write(c) for c in z])
+            f.close()
+        except Exception:
+            raise
+        finally:
+            os.remove(f.name)
+
+    def test_write_iterable_zip64_with_allow_zip64_many_smalls(self):
+        # check many small streams that sum length require ZIP64 extensions when allowed zip64
+        z = zipstream.ZipFile(mode='w', allowZip64=True)
+
+        def string_small_generator():
+            counter = 0
+            sample = b'zipstream0' * 10000000
+            len_sample = len(sample)
+            while counter + len_sample < zipstream.ZIP64_LIMIT:
+                counter += len_sample
+                yield sample
+
+        data = [string_small_generator(), string_small_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        try:
+            for chunk in z:
+                f.write(chunk)
+            f.close()
+        except Exception:
+            raise
+        finally:
+            os.remove(f.name)
+
+    def test_write_iterable_zip64_with_allow_zip64_1_big_file(self):
+        # check 1 big stream that length require ZIP64 extensions when allowed zip64
+        z = zipstream.ZipFile(mode='w', allowZip64=True)
+
+        def string_big_generator():
+            counter = 0
+            sample = b'zipstream0' * 10000000
+            len_sample = len(sample)
+            while counter < zipstream.ZIP64_LIMIT:
+                counter += len_sample
+                yield sample
+
+        data = [string_big_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        try:
+            for chunk in z:
+                f.write(chunk)
+            f.close()
+        except Exception:
+            raise
+        finally:
+            os.remove(f.name)
+
+
 if __name__ == '__main__':
     unittest.main()

zipstream/__init__.py

@@ -324,6 +324,11 @@ def __write(self, filename=None, iterable=None, arcname=None, compress_type=None
             zinfo.compress_size = file_size
         zinfo.CRC = CRC
         zinfo.file_size = file_size
+        if not filename:
+            zip64 = zinfo.file_size > ZIP64_LIMIT or \
+                    zinfo.compress_size > ZIP64_LIMIT
+            if zip64 and not self._allowZip64:
+                raise zipfile.LargeZipFile("Filesize would require ZIP64 extensions")
         if not zip64 and self._allowZip64:
             if file_size > ZIP64_LIMIT:
                 raise RuntimeError('File size has increased during compressing')
@@ -421,6 +426,8 @@ def __close(self):
                     centDirOffset > ZIP64_LIMIT or
                     centDirSize > ZIP64_LIMIT):
                     # Need to write the ZIP64 end-of-archive records
+                    if not self._allowZip64:
+                        raise zipfile.LargeZipFile(" would require ZIP64 extensions")
                     zip64endrec = struct.pack(
                             structEndArchive64, stringEndArchive64,
                             44, 45, 45, 0, 0, centDirCount, centDirCount,