Merge pull request #5 from sebdiem/master

Added support for iterable input instead of file

Author: allanlei

README.markdown

@@ -7,17 +7,55 @@
 zipstream.py is a zip archive generator based on python 3.3's zipfile.py. It was created to
 generate a zip file generator for streaming (ie web apps). This is beneficial for when you 
 want to provide a downloadable archive of a large collection of regular files, which would be infeasible to
-generate the archive prior to downloading.
+generate the archive prior to downloading or of a very large file that you do not want to store entirely on disk or on memory.
 
 The archive is generated as an iterator of strings, which, when joined, form
 the zip archive. For example, the following code snippet would write a zip
 archive containing files from 'path' to a normal file:
 
 ```python
-zf = open('zipfile.zip', 'wb')
-for data in ZipStream(path):
-    zf.write(data)
-zf.close()
+import zipstream
+
+z = zipstream.ZipFile()
+z.write('path/to/files')
+
+with open('zipfile.zip', 'wb') as f:
+    for data in z:
+        f.write(data)
+```
+
+zipstream also allows to take as input a byte string iterable and to generate
+the archive as an iterator.
+This avoids storing large files on disk or in memory.
+To do so you could use something like this snippet:
+
+```python
+def iterable():
+    for _ in xrange(10):
+        yield b'this is a byte string\x01\n'
+
+z = zipstream.ZipFile()
+z.write_iter(iterable(), 'my_archive_iter')
+
+with open('zipfile.zip', 'wb') as f:
+    for data in z:
+        f.write(data)
+```
+
+Of course both approach can be combined:
+
+```python
+def iterable():
+    for _ in xrange(10):
+        yield b'this is a byte string\x01\n'
+
+z = zipstream.ZipFile()
+z.write('path/to/files', 'my_archive_files')
+z.write_iter(iterable(), 'my_archive_iter')
+
+with open('zipfile.zip', 'wb') as f:
+    for data in z:
+        f.write(data)
 ```
 
 Since recent versions of web.py support returning iterators of strings to be
@@ -31,10 +69,10 @@ def GET(self):
     web.header('Content-type' , 'application/zip')
     web.header('Content-Disposition', 'attachment; filename="%s"' % (
         zip_filename,))
-    return ZipStream(path)
+    return zipstream.ZipFile(path)
 ```
 
-If the zlib module is available, ZipStream can generate compressed zip
+If the zlib module is available, zipstream.ZipFile can generate compressed zip
 archives.
 
 ## Requirements
@@ -51,7 +89,7 @@ from flask import Response
 @app.route('/package.zip', methods=['GET'], endpoint='zipball')
 def zipball():
     def generator():
-    	z = ZipStream(mode='w', compression=ZIP_DEFLATED)
+    	z = zipstream.ZipFile(mode='w', compression=ZIP_DEFLATED)
 
     	z.write('/path/to/file')
 
@@ -66,7 +104,7 @@ def zipball():
 
 @app.route('/package.zip', methods=['GET'], endpoint='zipball')
 def zipball():
-	z = ZipStream(mode='w', compression=ZIP_DEFLATED)
+	z = zipstream.ZipFile(mode='w', compression=ZIP_DEFLATED)
 	z.write('/path/to/file')
 
     response = Response(z, mimetype='application/zip')
@@ -80,7 +118,7 @@ def zipball():
 from django.http import StreamingHttpResponse
 
 def zipball(request):
-	z = ZipStream(mode='w', compression=ZIP_DEFLATED)
+	z = zipstream.ZipFile(mode='w', compression=ZIP_DEFLATED)
 	z.write('/path/to/file')
 
     response = StreamingHttpResponse(z, mimetype='application/zip')
@@ -97,5 +135,11 @@ def GET(self):
     web.header('Content-type' , 'application/zip')
     web.header('Content-Disposition', 'attachment; filename="%s"' % (
         zip_filename,))
-    return ZipStream(path)
+    return zipstream.ZipFile(path)
 ```
+
+## Running tests
+    
+With python version > 2.6, just run the following command: `python -m unittest discover`
+
+Alternatively, you can use `nose`.

tests/__init__.py

@@ -50,10 +50,33 @@ def test_write_file(self):
         f.close()
 
         z2 = zipfile.ZipFile(f.name, 'r')
-        z2.testzip()
+        self.assertFalse(z2.testzip())
 
         os.remove(f.name)
 
+    def test_write_iterable(self):
+        z = zipstream.ZipFile(mode='w')
+        def string_generator():
+            for _ in range(10):
+                yield b'zipstream\x01\n'
+        data = [string_generator(), string_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        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))
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()

tests/test_zipstream.py

@@ -170,8 +170,8 @@ def __init__(self, fileobj=None, mode='w', compression=ZIP_STORED, allowZip64=Fa
         self.paths_to_write = []
 
     def __iter__(self):
-        for args, kwargs in self.paths_to_write:
-            for data in self.__write(*args, **kwargs):
+        for kwargs in self.paths_to_write:
+            for data in self.__write(**kwargs):
                 yield data
         for data in self.__close():
             yield data
@@ -204,21 +204,30 @@ def write(self, filename, arcname=None, compress_type=None):
         # TODO: Reflect python's Zipfile.write
         #   - if filename is file, write as file
         #   - if filename is directory, write an empty directory
-        self.paths_to_write.append(
-            ((filename, ), {'arcname': arcname, 'compress_type': compress_type}),
-        )
+        kwargs = {'filename': filename, 'arcname': arcname, 'compress_type': compress_type}
+        self.paths_to_write.append(kwargs)
 
-    def __write(self, filename, arcname=None, compress_type=None):
+    def write_iter(self, arcname, iterable, compress_type=None):
+        """Write the bytes iterable `iterable` to the archive under the name `arcname`."""
+        kwargs = {'arcname': arcname, 'iterable': iterable, 'compress_type': compress_type}
+        self.paths_to_write.append(kwargs)
+
+    def __write(self, filename=None, iterable=None, arcname=None, compress_type=None):
         """Put the bytes from filename into the archive under the name
-        arcname."""
+        `arcname`."""
         if not self.fp:
             raise RuntimeError(
                   "Attempt to write to ZIP archive that was already closed")
-
-        st = os.stat(filename)
-        isdir = stat.S_ISDIR(st.st_mode)
-        mtime = time.localtime(st.st_mtime)
-        date_time = mtime[0:6]
+        if (filename is None and iterable is None) or (filename is not None and iterable is not None):
+            raise ValueError("either (exclusively) filename or iterable shall be not None")
+
+        if filename:
+            st = os.stat(filename)
+            isdir = stat.S_ISDIR(st.st_mode)
+            mtime = time.localtime(st.st_mtime)
+            date_time = mtime[0:6]
+        else:
+            st, isdir, date_time = None, False, time.localtime()[0:6]
         # Create ZipInfo instance to store file information
         if arcname is None:
             arcname = filename
@@ -228,13 +237,16 @@ def __write(self, filename, arcname=None, compress_type=None):
         if isdir:
             arcname += '/'
         zinfo = ZipInfo(arcname, date_time)
-        zinfo.external_attr = (st[0] & 0xFFFF) << 16      # Unix attributes
+        if st:
+            zinfo.external_attr = (st[0] & 0xFFFF) << 16      # Unix attributes
+        else:
+            zinfo.external_attr = 0o600 << 16     # ?rw-------
         if compress_type is None:
             zinfo.compress_type = self.compression
         else:
             zinfo.compress_type = compress_type
 
-        zinfo.file_size = st.st_size
+        zinfo.file_size = 0
         zinfo.flag_bits = 0x00
         zinfo.flag_bits |= 0x08                 # ZIP flag bits, bit 3 indicates presence of data descriptor
         zinfo.header_offset = self.fp.tell()    # Start of header bytes
@@ -255,19 +267,29 @@ def __write(self, filename, arcname=None, compress_type=None):
             return
 
         cmpr = _get_compressor(zinfo.compress_type)
-        with open(filename, 'rb') as fp:
-            # Must overwrite CRC and sizes with correct data later
-            zinfo.CRC = CRC = 0
-            zinfo.compress_size = compress_size = 0
-            # Compressed size can be larger than uncompressed size
-            zip64 = self._allowZip64 and \
-                    zinfo.file_size * 1.05 > ZIP64_LIMIT
-            yield self.fp.write(zinfo.FileHeader(zip64))
-            file_size = 0
-            while 1:
-                buf = fp.read(1024 * 8)
-                if not buf:
-                    break
+
+        # Must overwrite CRC and sizes with correct data later
+        zinfo.CRC = CRC = 0
+        zinfo.compress_size = compress_size = 0
+        # Compressed size can be larger than uncompressed size
+        zip64 = self._allowZip64 and \
+                zinfo.file_size * 1.05 > ZIP64_LIMIT
+        yield self.fp.write(zinfo.FileHeader(zip64))
+        file_size = 0
+        if filename:
+            with open(filename, 'rb') as fp:
+                while 1:
+                    buf = fp.read(1024 * 8)
+                    if not buf:
+                        break
+                    file_size = file_size + len(buf)
+                    CRC = crc32(buf, CRC) & 0xffffffff
+                    if cmpr:
+                        buf = cmpr.compress(buf)
+                        compress_size = compress_size + len(buf)
+                    yield self.fp.write(buf)
+        else: # we have an iterable
+            for buf in iterable:
                 file_size = file_size + len(buf)
                 CRC = crc32(buf, CRC) & 0xffffffff
                 if cmpr: