From 4ea51e39b8f5f08c791820ba57ddf45764e76d56 Mon Sep 17 00:00:00 2001
From: Franklyn Tackitt <franklyn@tackitt.net>
Date: Wed, 15 May 2019 11:20:02 -0700
Subject: [PATCH] Add partial flushing of ZipStreams

I use this to flush partial zips as files are streamed into
 them from requests, and then at the end add manifest and
 error files to the end of the archive

I've also added a related test and example of use
---
 README.md               | 21 +++++++++++++++++++++
 tests/test_zipstream.py | 25 +++++++++++++++++++++++++
 tox.ini                 |  2 +-
 zipstream/__init__.py   | 11 ++++++++---
 4 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 994ec53..5c92a91 100644
--- a/README.md
+++ b/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+
diff --git a/tests/test_zipstream.py b/tests/test_zipstream.py
index 9910fe2..6bda736 100644
--- a/tests/test_zipstream.py
+++ b/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))
diff --git a/tox.ini b/tox.ini
index 8302cc5..d93b9c8 100644
--- a/tox.ini
+++ b/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
diff --git a/zipstream/__init__.py b/zipstream/__init__.py
index b824d76..5d7accd 100644
--- a/zipstream/__init__.py
+++ b/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."""