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."""