improve add_file_from_memory

closes #68

Author: Changaco

libarchive/write.py

@@ -86,8 +86,8 @@ def add_file_from_memory(
         :type entry_path: str
         :param entry_size: entire size of entry in bytes
         :type entry_size: int
-        :param entry_data: content of entry as a list of byte-like objects
-        :type entry_data: iterable
+        :param entry_data: content of entry
+        :type entry_data: bytes or Iterable[bytes]
         :param filetype: which type of file: normal, symlink etc.
         should entry be created as
         :type filetype: octal number
@@ -96,6 +96,13 @@ def add_file_from_memory(
         """
         archive_pointer = self._pointer
 
+        if isinstance(entry_data, bytes):
+            entry_data = (entry_data,)
+        elif isinstance(entry_data, str):
+            raise TypeError(
+                "entry_data: expected bytes, got %r" % type(entry_data)
+            )
+
         with new_archive_entry() as archive_entry_pointer:
             archive_entry = ArchiveEntry(None, archive_entry_pointer)
 

tests/test_rwx.py

@@ -126,11 +126,8 @@ def test_write_not_fail(write_fail_mock):
 def test_adding_entry_from_memory(archfmt, data_bytes):
     entry_path = 'testfile.data'
 
-    # entry_data must be a list of byte-like objects for maximum compatibility
-    # (Use of other data types can have undesirable side-effects)
-    # entry_size is the total size in bytes of the entry_data items
-    entry_data = [data_bytes]
-    entry_size = sum(len(bdata) for bdata in entry_data)
+    entry_data = data_bytes
+    entry_size = len(data_bytes)
 
     blocks = []
 
@@ -144,7 +141,7 @@ def write_callback(data):
     buf = b''.join(blocks)
     with libarchive.memory_reader(buf) as memory_archive:
         for archive_entry in memory_archive:
-            expected = b''.join(entry_data)
+            expected = entry_data
             actual = b''.join(archive_entry.get_blocks())
             assert expected == actual
             assert archive_entry.path == entry_path