Prevent mismatches between arcname and data when adding paths (#22)

Previously, it was possible to add a directory and give it an arcname that
didn't end with a "/". Likewise, it was also possible to add a file
with a arcname that ended with a "/". These cases would cause an
incorrect stream size to be calculated, as well as create strange zip
files.

This commit fixes this issue by modifying arcnames as paths are added to
ensure that directories end with "/" and files do not.

tests/test_zipstream.py

@@ -517,6 +517,31 @@ def test_directly_adding_empty_dir(tmpdir):
     assert zinfos[0].compress_size == 0
 
 
+def test_add_path_dir_as_file(tmpdir):
+    t = tmpdir.mkdir("dir")
+    zs = ZipStream(sized=True)
+    zs.add_path(
+        t,
+        arcname="filename.txt",
+        recurse=False
+    )
+    assert zs.info_list()[0]["name"] == "filename.txt/"
+    assert len(zs) == len(bytes(zs))
+
+
+def test_add_path_file_as_dir(tmpdir):
+    t = tmpdir.join("file")
+    t.write(b"x")
+    zs = ZipStream(sized=True)
+    zs.add_path(
+        t,
+        arcname="dir/",
+        recurse=False
+    )
+    assert zs.info_list()[0]["name"] == "dir"
+    assert len(zs) == len(bytes(zs))
+
+
 def test_adding_local_dir(tmpdir):
     """Test adding files/folder from the local directory doesn't include the directory name of ."""
     t = tmpdir.mkdir("top")

zipstream/ng.py

@@ -713,21 +713,9 @@ def add_path(self, path, arcname=None, *, recurse=True, compress_type=None, comp
         for filepath in recurse(path):
             filename = os.path.relpath(filepath, path)
             filearcname = os.path.normpath(os.path.join(arcname, filename))
-
-            # Check if adding a directory, and if so, add a trailing slash
-            # (normpath will remove it). Also set the size since we're doing
-            # the stat anyway
-            st = os.stat(filepath)
-            if stat.S_ISDIR(st.st_mode):
-                filearcname += "/"
-                size = 0
-            else:
-                size = st.st_size
-
             self._enqueue(
                 path=filepath,
                 arcname=filearcname,
-                size=size,
                 compress_type=compress_type,
                 compress_level=compress_level
             )
@@ -993,15 +981,25 @@ 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:
+        # Get the expected size of the data where possible
+        if path or size is None:
+            if path:
                 if stat.S_ISDIR(st.st_mode):
                     kwargs["size"] = 0
                 else:
                     kwargs["size"] = st.st_size
+            elif data is not None:
+                kwargs["size"] = len(data)
+
+        # If adding data from a path ensure that the name matches the type of
+        # data being added.
+        if path:
+            name_is_dir = kwargs["arcname"][-1] == "/"
+            path_is_dir = stat.S_ISDIR(st.st_mode)
+            if path_is_dir and not name_is_dir:
+                kwargs["arcname"] += "/"
+            elif not path_is_dir and name_is_dir:
+                kwargs["arcname"] = kwargs["arcname"].rstrip("/")
 
         # If the ZipStream is sized then it will look at what is being added
         # and add the number of bytes used by this file to the running total