Enhance add_file_from_memory

libarchive/write.py

@@ -84,10 +84,10 @@ def add_file_from_memory(
 
         :param entry_path: where entry should be places in archive
         :type entry_path: str
-        :param entry_size: entire size of entry
+        :param entry_size: entire size of entry in bytes
         :type entry_size: int
         :param entry_data: content of entry
-        :type entry_data: iterable
+        :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

@@ -1,12 +1,15 @@
 """Test reading, writing and extracting archives."""
 
 from __future__ import division, print_function, unicode_literals
+
 import io
+import json
 
 import libarchive
 from libarchive.extract import EXTRACT_OWNER, EXTRACT_PERM, EXTRACT_TIME
 from libarchive.write import memory_writer
 from mock import patch
+import pytest
 
 from . import check_archive, in_dir, treestat
 
@@ -115,24 +118,30 @@ def test_write_not_fail(write_fail_mock):
     assert not write_fail_mock.called
 
 
-def test_adding_entry_from_memory():
-    entry_path = 'this is path'
-    entry_data = 'content'
-    entry_size = len(entry_data)
+@pytest.mark.parametrize(
+    'archfmt,data_bytes',
+    [('zip', b'content'),
+     ('gnutar', b''),
+     ('pax', json.dumps({'a': 1, 'b': 2, 'c': 3}).encode()),
+     ('7zip', b'lorem\0ipsum')])
+def test_adding_entry_from_memory(archfmt, data_bytes):
+    entry_path = 'testfile.data'
+    entry_data = data_bytes
+    entry_size = len(data_bytes)
 
     blocks = []
 
     def write_callback(data):
         blocks.append(data[:])
         return len(data)
 
-    with libarchive.custom_writer(write_callback, 'zip') as archive:
+    with libarchive.custom_writer(write_callback, archfmt) as archive:
         archive.add_file_from_memory(entry_path, entry_size, entry_data)
 
     buf = b''.join(blocks)
     with libarchive.memory_reader(buf) as memory_archive:
         for archive_entry in memory_archive:
-            assert entry_data.encode() == b''.join(
-                archive_entry.get_blocks()
-            )
+            expected = entry_data
+            actual = b''.join(archive_entry.get_blocks())
+            assert expected == actual
             assert archive_entry.path == entry_path