Merge pull request #1 from DisruptiveLabs/partial-flush

Add partial flushing of ZipStreams

Author: arjan-s

README.md

@@ -113,6 +113,27 @@ def zipball():
     response = Response(z, mimetype='application/zip')
     response.headers['Content-Disposition'] = 'attachment; filename={}'.format('files.zip')
     return response
+
+# Partial flushing of the zip before closing
+
+@app.route('/package.zip', methods=['GET'], endpoint='zipball')
+def zipball():
+    def generate_zip_with_manifest():
+        z = zipstream.ZipFile(mode='w', compression=ZIP_DEFLATED)
+
+        manifest = []
+        for filename in os.listdir('/path/to/files'):
+            z.write(os.path.join('/path/to/files', filename), arcname=filename)
+            yield from z.flush()
+            manifest.append(filename)
+
+        z.write_str('manifest.json', json.dumps(manifest).encode())
+
+        yield from z
+
+    response = Response(z, mimetype='application/zip')
+    response.headers['Content-Disposition'] = 'attachment; filename={}'.format('files.zip')
+    return response
 ```
 
 ### django 1.5+

tests/test_zipstream.py

@@ -92,6 +92,31 @@ def test_writestr(self):
 
         os.remove(f.name)
 
+    def test_partial_writes(self):
+        z = zipstream.ZipFile(mode='w')
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+
+        with open(SAMPLE_FILE_RTF, 'rb') as fp:
+            z.writestr('sample1.rtf', fp.read())
+
+        for chunk in z.flush():
+            f.write(chunk)
+
+        with open(SAMPLE_FILE_RTF, 'rb') as fp:
+            z.writestr('sample2.rtf', fp.read())
+
+        for chunk in z.flush():
+            f.write(chunk)
+
+        for chunk in z:
+            f.write(chunk)
+        
+        f.close()
+        z2 = zipfile.ZipFile(f.name, 'r')
+        self.assertFalse(z2.testzip())
+
+        os.remove(f.name)
+
     def test_write_iterable_no_archive(self):
         z = zipstream.ZipFile(mode='w')
         self.assertRaises(TypeError, z.write_iter, iterable=range(10))

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py26, py27, py32, py33, py34, py35, pypy
+envlist = py26, py27, py32, py33, py34, py35, py36, py37, pypy
 
 [testenv]
 deps=nose

zipstream/__init__.py

@@ -178,9 +178,8 @@ def __init__(self, fileobj=None, mode='w', compression=ZIP_STORED, allowZip64=Fa
         self.paths_to_write = []
 
     def __iter__(self):
-        for kwargs in self.paths_to_write:
-            for data in self.__write(**kwargs):
-                yield data
+        for data in self.flush():
+            yield data
         for data in self.__close():
             yield data
 
@@ -190,6 +189,12 @@ def __enter__(self):
     def __exit__(self, type, value, traceback):
         self.close()
 
+    def flush(self):
+        while self.paths_to_write:
+            kwargs = self.paths_to_write.pop()
+            for data in self.__write(**kwargs):
+                yield data
+
     @property
     def comment(self):
         """The comment text associated with the ZIP file."""