Merge pull request #176 from uktrade/feat/better-buffering

feat: remove input buffering and flush output buffer after local header

Author: michalc

docs/api/functions.md

@@ -33,7 +33,7 @@ def stream_zip(
 | Name                | Type                           | Description
 | --------------------| -------------------------------| ------------------------------------------
 | files               | Iterable[MemberFile]           | The member files of the ZIP
-| chunk_size          | int                            | How many bytes of each member file to fetch before compressing
+| chunk_size          | int                            | Maximum `bytes` instance length yielded to client code
 | get_compressobj     | Callable[[], 'zlib._Compress'] | A function returning a [Python zlib compression object](https://docs.python.org/3/library/zlib.html#zlib.compressobj) |
 | password            | Optional[str]                  | The password used to encrypt all the member files with AES-256 encryption adhering to the Winzip AE-2 specification - see [Password protection](/get-started/password-protection/)
 | extended_timestamps | bool                           | Whether to save extended timestamps in the ZIP file
@@ -82,7 +82,7 @@ async def async_stream_zip(
 | Name                | Type                           | Description
 | --------------------| -------------------------------| ------------------------------------------
 | files               | AsyncIterable[AsyncMemberFile] | The member files of the ZIP
-| chunk_size          | int                            | How many bytes of each member file to fetch before compressing
+| chunk_size          | int                            | Maximum `bytes` instance length yielded to client code
 | get_compressobj     | Callable[[], 'zlib._Compress'] | A function returning a [Python zlib compression object](https://docs.python.org/3/library/zlib.html#zlib.compressobj) |
 | password            | Optional[str]                  | The password used to encrypt all the member files with AES-256 encryption adhering to the Winzip AE-2 specification - see [Password protection](/get-started/password-protection/)
 | extended_timestamps | bool                           | Whether to save extended timestamps in the ZIP file

docs/get-started/advanced-usage.md

@@ -23,16 +23,16 @@ If you wish to disable compression entirely for these methods, you can pass `lev
 
 ## Custom chunk size
 
-The default `bytes` instance size is 65536 bytes. To customise this, you can override the `chunk_size` parameter.
+The default `bytes` instance size for output is 65536 bytes. To customise this, you can override the `chunk_size` parameter.
 
 ```python
 for zipped_chunk in stream_zip(unzipped_files(), chunk_size=65536):
     print(zipped_chunk)
 ```
 
-This one size is used both for input - splitting or gathering any uncompressed data into `chunk_size` bytes before attempting to compress it, and in output - splitting or gathering any compressed data into `chunk_size` bytes before returning it to client code.
+This size is used to buffer output - splitting or gathering any compressed data into a `chunk_size` buffer before yielding it to client code. There may be performance differences with different `chunk_size` values; the default chunk_size may not be optimal for your use case.
 
-There may be performance differences with different `chunk_size` values. The default chunk_size may not be optimal for your use case.
+The buffer is flushed just before iterating over the bytes of each member file, irrespective of `chunk_size`.
 
 
 ## Extended timestamps

stream_zip/__init__.py

@@ -96,6 +96,13 @@ def _get(self, offset: int, default_get_compressobj: _CompressObjGetter) -> _Met
 
         return _ZIP_AUTO_TYPE_INNER()
 
+# Sentinal object as a command that output buffer be flushed
+# Extends from bytes to pass type checking
+class _Flusher(bytes):
+    pass
+
+_flush = _Flusher()
+
 
 ###############################
 # Public sentinel objects/types
@@ -135,6 +142,9 @@ def up_to(num: int) -> Iterable[bytes]:
                         break
                     else:
                         offset = 0
+                        if chunk is _flush:
+                            chunk = b''
+                            break
                 to_yield = min(num, len(chunk) - offset)
                 offset = offset + to_yield
                 num -= to_yield
@@ -280,6 +290,7 @@ def _zip_64_local_header_and_data(
             ))
             yield from _(name_encoded)
             yield from _(extra)
+            yield _flush
 
             uncompressed_size, raw_compressed_size, crc_32 = yield from encryption_func(_zip_data(
                 chunks,
@@ -347,6 +358,7 @@ def _zip_32_local_header_and_data(
             ))
             yield from _(name_encoded)
             yield from _(extra)
+            yield _flush
 
             uncompressed_size, raw_compressed_size, crc_32 = yield from encryption_func(_zip_data(
                 chunks,
@@ -444,6 +456,7 @@ def _no_compression_64_local_header_and_data(
             ))
             yield from _(name_encoded)
             yield from _(extra)
+            yield _flush
 
             yield from encryption_func((chunk for chunk in chunks))
 
@@ -506,6 +519,7 @@ def _no_compression_32_local_header_and_data(
             ))
             yield from _(name_encoded)
             yield from _(extra)
+            yield _flush
 
             yield from encryption_func((chunk for chunk in chunks))
 
@@ -582,6 +596,7 @@ def _no_compression_streamed_64_local_header_and_data(
             ))
             yield from _(name_encoded)
             yield from _(extra)
+            yield _flush
 
             yield from encryption_func(_no_compression_streamed_data(chunks, uncompressed_size, crc_32, 0xffffffffffffffff))
 
@@ -642,6 +657,7 @@ def _no_compression_streamed_32_local_header_and_data(
             ))
             yield from _(name_encoded)
             yield from _(extra)
+            yield _flush
 
             yield from encryption_func(_no_compression_streamed_data(chunks, uncompressed_size, crc_32, 0xffffffff))
 
@@ -716,7 +732,7 @@ def _no_compression_streamed_data(chunks: Iterable[bytes], uncompressed_size: in
                 (99, 28, aes_flag, aes_extra_struct.pack(aes_extra_signature, 7, 2, b'AE', 3, raw_compression), 0, _get_encrypt_aes(password)) if password is not None else \
                 (raw_compression, 0, 0, b'', 0xffffffff, _encrypt_dummy)
 
-            central_directory_header_entry, name_encoded, extra = yield from data_func(compression, aes_size_increase, aes_flags, name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, crc_32_mask, _get_compress_obj, encryption_func, evenly_sized(chunks))
+            central_directory_header_entry, name_encoded, extra = yield from data_func(compression, aes_size_increase, aes_flags, name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, crc_32_mask, _get_compress_obj, encryption_func, chunks)
             central_directory_size += len(central_directory_header_signature) + len(central_directory_header_entry) + len(name_encoded) + len(extra)
             central_directory.append((central_directory_header_entry, name_encoded, extra))
 
@@ -786,8 +802,7 @@ def _no_compression_streamed_data(chunks: Iterable[bytes], uncompressed_size: in
                 0, # ZIP_32 file comment length
             ))
 
-    zipped_chunks = get_zipped_chunks_uneven()
-    yield from evenly_sized(zipped_chunks)
+    yield from evenly_sized(get_zipped_chunks_uneven())
 
 
 async def async_stream_zip(

test_stream_zip.py

@@ -1,3 +1,4 @@
+from collections import Counter
 from datetime import datetime, timezone, timedelta
 from io import BytesIO
 import asyncio
@@ -152,7 +153,7 @@ def files():
         yield 'file-1', now, mode, method(1, zlib.crc32(b'')), (b'a',)
 
     with pytest.raises(CRC32IntegrityError):
-        for name, size, chunks in stream_unzip(stream_zip(files())):
+        for _ in stream_zip(files()):
             pass
 
 
@@ -173,7 +174,8 @@ def files():
 
     with pytest.raises(UncompressedSizeIntegrityError):
         for name, size, chunks in stream_unzip(stream_zip(files())):
-            pass
+            for _ in chunks:
+                pass
 
 
 def test_with_stream_unzip_auto_small():
@@ -762,10 +764,71 @@ def get_sizes():
             yield len(chunk)
 
     sizes = list(get_sizes())
-    assert set(sizes[:-1]) == {65536}
+    assert sizes[0] == 65
+    assert set(sizes[1:-1]) == {65536}
     assert sizes[-1] <= 65536
 
 
+@pytest.mark.parametrize(
+    "method",
+    [
+        ZIP_32,
+        ZIP_64,
+        NO_COMPRESSION_64(1, 2547889144),
+        NO_COMPRESSION_32(1, 2547889144),
+    ],
+)
+def test_local_headers_flushed(method):
+    now = datetime.strptime('2021-01-01 21:01:12', '%Y-%m-%d %H:%M:%S')
+    mode = stat.S_IFREG | 0o600
+
+    state = []
+
+    def data():
+        state.append('data')
+        yield b'-'
+
+    def files():
+        state.append('file')
+        yield 'file-1', now, mode, method, data()
+
+    for chunk in stream_zip(files()):
+        state.append('chunk')
+        if b'file-1' in chunk:
+            state.append('file-name')
+
+    assert state == ['file', 'chunk', 'file-name', 'data', 'chunk', 'file-name']
+
+
+@pytest.mark.parametrize(
+    "method",
+    [
+        NO_COMPRESSION_64,
+        NO_COMPRESSION_32,
+    ],
+)
+def test_local_headers_flushed_buffered_data(method):
+    now = datetime.strptime('2021-01-01 21:01:12', '%Y-%m-%d %H:%M:%S')
+    mode = stat.S_IFREG | 0o600
+
+    state = []
+
+    def data():
+        state.append('data')
+        yield b'-'
+
+    def files():
+        state.append('file')
+        yield 'file-1', now, mode, method, data()
+
+    for chunk in stream_zip(files()):
+        state.append('chunk')
+        if b'file-1' in chunk:
+            state.append('file-name')
+
+    assert state == ['file', 'data', 'chunk', 'file-name', 'chunk', 'file-name']
+
+
 @pytest.mark.parametrize(
     "method",
     [
@@ -1384,7 +1447,7 @@ async def test():
             state.append('out')
 
     asyncio.get_event_loop().run_until_complete(test())
-    assert state == ['in', 'in', 'out', 'in', 'out', 'in', 'out', 'out']
+    assert state == ['out', 'in', 'in', 'out', 'in', 'out', 'in', 'out', 'out']
 
 
 @pytest.mark.skipif(