Repository: pR0Ps/zipstream-ng (language: python)

Issue:
Allow specifying the size of iterables added to the stream
Currently, if a `ZipStream` is sized (`ZipStream(sized=True)`) and an iterable is added, all the data in the iterable is pulled into memory to ensure the total size is known: https://github.com/pR0Ps/zipstream-ng/blob/16cfd3de814084b8af65617b649efde88472ac3d/zipstream/ng.py#L769-L773

In cases where the total size of the data from the iterable is known ahead of time, it should be possible to pass it into the library instead. An example of where this would be useful is for data downloaded from an HTTP server that provides a `Content-Length` header. The data should be streamed, but the total size is known (see #11)

Since getting the size wrong will lead to an inaccurate length computation, the passed in length should be validated after the iterator is run (this is also probably a good idea in general - what if the size of a file changes after being added but before being streamed?)


Discussion / hints:

Fix patch:
diff --git a/zipstream/ng.py b/zipstream/ng.py
index e5e548d..cd9a71f 100644
--- a/zipstream/ng.py
+++ b/zipstream/ng.py
@@ -7,7 +7,6 @@
 import functools
 import logging
 import os
-import queue
 import stat
 import struct
 import sys
@@ -423,7 +422,7 @@ def _iter_file(path):
             yield buf
 
 
-class ZipStream(object):
+class ZipStream:
     """A write-only zip that is generated from source files/data as it's
     iterated over.
 
@@ -458,9 +457,10 @@ def __init__(self, *, compress_type=ZIP_STORED, compress_level=None, sized=False
             prior to being generated, making it work with the `len()` function.
             Enabling this will enforce two restrictions:
               - No compression can be used
-              - Any iterables added to the stream will immediately be read fully
-                into memory since the size of them needs to be known at the time
-                they are added.
+              - Any iterables added to the stream without also specifying their
+                size (see `.add` docs) will immediately be read fully into
+                memory. This is because the size of the data they will produce
+                must be known prior to the stream being generated.
 
             If `False` (the default), no restrictions are enforced and using the
             object with the `len()` function will not work (will raise a
@@ -649,7 +649,7 @@ def _getkey(path):
 
     @_validate_final
     @_validate_compression
-    def add(self, data, arcname, *, compress_type=None, compress_level=None):
+    def add(self, data, arcname, *, size=None, compress_type=None, compress_level=None):
         """Queue up data to be added to the ZipStream
 
         `data` can be bytes, a string (encoded to bytes using utf-8), or any
@@ -663,6 +663,10 @@ def add(self, data, arcname, *, compress_type=None, compress_level=None):
         any `data` is provided then the `arcname` cannot end with a "/" as this
         would create a directory (which can't contain content).
 
+        `size` (optional) specifies the size of the `data` ONLY in the case
+        where it is an iterator and the ZipStream is sized. It is ignored in
+        all other cases.
+
         Note that the data provided will not be used until the file is actually
         encoded in the ZipStream. This means that strings and bytes will be held
         in memory and iterables will not be iterated over until then. For this
@@ -701,6 +705,7 @@ def add(self, data, arcname, *, compress_type=None, compress_level=None):
         elif hasattr(data, "__iter__"):
             self._enqueue(
                 iterable=data,
+                size=size if self._sized else None,
                 arcname=arcname,
                 compress_type=compress_type,
                 compress_level=compress_level,
@@ -819,6 +824,8 @@ def _enqueue(self, **kwargs):
         """Internal method to enqueue files, data, and iterables to be streamed"""
 
         path = kwargs.get("path")
+        data = kwargs.get("data")
+        size = kwargs.get("size")
 
         # Get the modified time of the added path (use current time for
         # non-paths) and use it to update the last_modified property
@@ -826,29 +833,31 @@ def _enqueue(self, **kwargs):
         if self._last_modified is None or self._last_modified < mtime:
             self._last_modified = mtime
 
+        # Get the expected size of the data where not specified and possible
+        if size is None:
+            if data is not None:
+                kwargs["size"] = len(data)
+            elif path is not None:
+                if os.path.isdir(path):
+                    kwargs["size"] = 0
+                else:
+                    kwargs["size"] = os.path.getsize(path)
+
         # If the ZipStream is sized then it will look at what is being added and
         # queue up some information for _get_size to use to compute the total
         # length of the stream. It will also read any iterables fully into
         # memory so their size is known.
         if self._sized:
-            iterable = kwargs.pop("iterable", None)
-            data = kwargs.get("data")
-
-            assert bool(iterable) ^ bool(data is not None) ^ bool(path)
-
             if kwargs.get("compress_type"):
                 raise ValueError("Cannot use compression with a sized ZipStream")
-            elif data is not None:
-                filesize = len(data)
-            elif iterable:
-                # Iterate the data to get the size and replace it with the static data
-                data = b''.join(iterable)
-                filesize = len(data)
+
+            # Iterate the iterable data to get the size and replace it with the static data
+            if path is None and data is None and size is None:
+                data = b''.join(kwargs.pop("iterable"))
+                kwargs["size"] = len(data)
                 kwargs["data"] = data
-            elif path:
-                filesize = os.path.getsize(path)
 
-            self._to_count.append((kwargs["arcname"], filesize))
+            self._to_count.append((kwargs["arcname"], kwargs["size"]))
 
         self._queue.append(kwargs)
 
@@ -857,9 +866,10 @@ def _track(self, data):
         self._pos += len(data)
         return data
 
-    def _gen_file_entry(self, *, path=None, iterable=None, data=None, arcname, compress_type, compress_level):
+    def _gen_file_entry(self, *, path=None, iterable=None, data=None, size=None, arcname, compress_type, compress_level):
         """Yield the zipped data generated by the specified path/iterator/data"""
         assert bool(path) ^ bool(iterable) ^ bool(data is not None)
+        assert not (self._sized and size is None)
 
         if path:
             zinfo = ZipStreamInfo.from_file(path, arcname)
@@ -900,6 +910,25 @@ def gen():
         # Generate the file data
         for x in zinfo._file_data(iterable, force_zip64=force_zip64):
             yield self._track(x)
+
+        if size is not None and size != zinfo.file_size:
+            # The size of the data that was stored didn't match what was
+            # expected. Note that this still produces a valid zip file, just
+            # one with a different amount of data than was expected.
+            # If the ZipStream is sized, this will raise an error since the
+            # actual size will no longer match the calculated size.
+            __log__.warning(
+                "Size mismatch when adding data for '%s' (expected %d bytes, got %d)",
+                arcname,
+                size,
+                zinfo.file_size
+            )
+            if self._sized:
+                raise RuntimeError(
+                    "Error adding '{}' to sized ZipStream - "
+                    "actual size did not match the computed size".format(arcname)
+                )
+
         self._filelist.append(zinfo)
 
     def _gen_archive_footer(self):


Test patch:
diff --git a/tests/test_zipstream.py b/tests/test_zipstream.py
index 0fd8c8a..3631b88 100644
--- a/tests/test_zipstream.py
+++ b/tests/test_zipstream.py
@@ -115,13 +115,18 @@ def _gen_rand():
     zipfile.ZIP_DEFLATED,
     zipfile.ZIP_BZIP2
 ])
-def test_zipstream_compression(files, ct):
+def test_zipstream_compression(caplog, files, ct):
     """Test that all types of compression properly compress and extract"""
+    caplog.set_level(logging.WARNING)
+
     zs = ZipStream(compress_type=ct)
     for f in files:
         zs.add_path(f[0])
 
     zf = _get_zip(zs)
+
+    assert "Size mismatch" not in caplog.text
+
     zinfos = zf.infolist()
     assert len(zinfos) == len(files)
 
@@ -498,9 +503,17 @@ def test_recursion_disable(tmpdir):
     [b'a', b'list', b'of', b'bytes'],
     _gen_rand()
 ])
-def test_adding_data(data):
-    """Test adding non-files"""
-    zs = ZipStream()
+@pytest.mark.parametrize("ct", [
+    zipfile.ZIP_STORED,
+    zipfile.ZIP_LZMA,
+    zipfile.ZIP_DEFLATED,
+    zipfile.ZIP_BZIP2
+])
+def test_adding_data(caplog, data, ct):
+    """Test adding non-files with different compression methods"""
+    caplog.set_level(logging.WARNING)
+
+    zs = ZipStream(compress_type=ct)
 
     tostore = data
     if isinstance(data, str):
@@ -519,9 +532,13 @@ def test_adding_data(data):
     zs.add(tostore, "data.bin")
     zf = _get_zip(zs)
 
+    assert "Size mismatch" not in caplog.text
+
     zinfos = zf.infolist()
     assert len(zinfos) == 1
 
+    assert zinfos[0].compress_type == ct
+
     from_file = zf.read(zinfos[0])
     assert from_file == data
 
@@ -590,6 +607,68 @@ def test_adding_missing_path(tmpdir):
         ZipStream.from_path(tmpdir.join("doesntexist"))
 
 
+def test_adding_deleted_path(tmpdir):
+    zs = ZipStream(sized=True)
+    f = tmpdir.join("file")
+    f.write(b"this is some data")
+    zs.add_path(f)
+    assert len(zs) == 139
+    f.remove()
+    assert len(zs) == 139
+
+    with pytest.raises(FileNotFoundError):
+        bytes(zs)
+
+
+@pytest.mark.parametrize("data", [
+    ("short text", "then longer text"),
+    ("first longer text", "then shorter"),
+])
+@pytest.mark.parametrize("sized", (True, False))
+def test_adding_changed_path(caplog, tmpdir, sized, data):
+    caplog.set_level(logging.WARNING)
+
+    zs = ZipStream(sized=sized)
+    f = tmpdir.join("file")
+    f.write(data[0])
+    zs.add_path(f)
+    f.write(data[1])  # overwrites, doesn't append
+
+    assert not caplog.text
+    if sized:
+        with pytest.raises(RuntimeError, match="Error adding 'file' to sized ZipStream"):
+            bytes(zs)
+    else:
+        bytes(zs)
+
+    msg = "Size mismatch when adding data for 'file' (expected {} bytes, got {})".format(
+        len(data[0]),
+        len(data[1])
+    )
+    assert msg in caplog.text
+
+
+@pytest.mark.parametrize("sized", (True, False))
+@pytest.mark.parametrize("size", (None, -1, 0, 1, 100))
+def test_adding_iterator_invalid_size(caplog, sized, size):
+    """Make sure errors are raised when generating if an incorrect iterator size was provided"""
+    caplog.set_level(logging.WARNING)
+
+    l_rand = len(b''.join(_gen_rand()))
+
+    zs = ZipStream(sized=sized)
+    zs.add(_gen_rand(), "rand.txt", size=size)
+
+    assert not caplog.text
+    if sized and size is not None:
+        with pytest.raises(RuntimeError, match="Error adding 'rand.txt' to sized ZipStream"):
+            bytes(zs)
+        assert "Size mismatch when adding data for 'rand.txt'" in caplog.text
+    else:
+        bytes(zs)
+        assert "Size mismatch" not in caplog.text
+
+
 def test_adding_comment(caplog):
     caplog.set_level(logging.WARNING)
 
@@ -904,6 +983,7 @@ def generate_zipstream(path):
         _get_zip(ZipStream.from_path(t))
     )
 
+
 def test_add_duplicate_file():
     """Test adding multiple files with the same name works"""
     zs = ZipStream(sized=True)
@@ -924,7 +1004,7 @@ def test_add_duplicate_file():
 
 
 def test_unsized_zipstream_len_typeerror():
-    """Test that an unsized ZipStream raises a TypeError when aske for the length"""
+    """Test that an unsized ZipStream raises a TypeError when asked for the length"""
     zs = ZipStream(sized=False)
     assert not zs.sized
 
@@ -934,7 +1014,7 @@ def test_unsized_zipstream_len_typeerror():
 
 @pytest.mark.parametrize("zip64", [False, True])
 @pytest.mark.parametrize("sized", [False, True])
-def test_proper_zip64_min_version(monkeypatch, files, zip64, sized):
+def test_proper_zip64_min_version(monkeypatch, zip64, sized):
     """Ensure that the min version is set properly"""
     if zip64:
         monkeypatch.setattr(zipfile, "ZIP64_LIMIT", 100)
@@ -954,6 +1034,27 @@ def test_proper_zip64_min_version(monkeypatch, files, zip64, sized):
     assert zi.extract_version == expected_version
 
 
+@pytest.mark.parametrize("sized", [False, True])
+@pytest.mark.parametrize("known_size", [False, True])
+def test_adding_iterator_sized(sized, known_size):
+    """Check that adding an iterator to the ZipStream only immediately reads it
+    into memory if the ZipStream is sized and the size isn't provided
+    """
+    read = False
+    def _my_iter():
+        nonlocal read
+        yield from _gen_rand()
+        read = True
+
+    l_rand = len(b''.join(_gen_rand()))
+
+    zs = ZipStream(sized=sized)
+    zs.add(_my_iter(), "rand.txt", size=None if not known_size else l_rand)
+    assert read == (False if known_size else sized)
+    bytes(zs)
+    assert read
+
+
 @pytest.mark.parametrize("zip64", [False, True])
 def test_sized_zipstream(monkeypatch, files, zip64):
     """Test a sized ZipStream accurately calculates its final size"""