Merge pull request #60 from versity/nic/custom-reader-clean

Add custom_reader for archives

Author: Changaco

libarchive/__init__.py

@@ -1,13 +1,13 @@
 from .entry import ArchiveEntry
 from .exception import ArchiveError
 from .extract import extract_fd, extract_file, extract_memory
-from .read import fd_reader, file_reader, memory_reader
+from .read import custom_reader, fd_reader, file_reader, memory_reader, stream_reader
 from .write import custom_writer, fd_writer, file_writer, memory_writer
 
 __all__ = [
     ArchiveEntry,
     ArchiveError,
     extract_fd, extract_file, extract_memory,
-    fd_reader, file_reader, memory_reader,
+    custom_reader, fd_reader, file_reader, memory_reader, stream_reader,
     custom_writer, fd_writer, file_writer, memory_writer
 ]

libarchive/ffi.py

@@ -44,6 +44,9 @@
 WRITE_CALLBACK = CFUNCTYPE(
     c_ssize_t, c_void_p, c_void_p, POINTER(c_void_p), c_size_t
 )
+READ_CALLBACK = CFUNCTYPE(
+    c_ssize_t, c_void_p, c_void_p, POINTER(c_void_p)
+)
 OPEN_CALLBACK = CFUNCTYPE(c_int, c_void_p, c_void_p)
 CLOSE_CALLBACK = CFUNCTYPE(c_int, c_void_p, c_void_p)
 VOID_CB = lambda *_: ARCHIVE_OK
@@ -170,6 +173,9 @@ def ffi(name, argtypes, restype, errcheck=None):
         logger.warning('read filter "%s" is not supported' % f_name)
         READ_FILTERS.remove(f_name)
 
+ffi('read_open',
+    [c_archive_p, c_void_p, OPEN_CALLBACK, READ_CALLBACK, CLOSE_CALLBACK],
+    c_int, check_int)
 ffi('read_open_fd', [c_archive_p, c_int, c_size_t], c_int, check_int)
 ffi('read_open_filename_w', [c_archive_p, c_wchar_p, c_size_t],
     c_int, check_int)

libarchive/read.py

@@ -1,11 +1,12 @@
 from __future__ import division, print_function, unicode_literals
 
 from contextlib import contextmanager
-from ctypes import cast, c_void_p
+from ctypes import cast, c_void_p, POINTER, create_string_buffer
 from os import fstat, stat
 
 from . import ffi
-from .ffi import ARCHIVE_EOF
+from .ffi import (ARCHIVE_EOF, OPEN_CALLBACK, READ_CALLBACK, CLOSE_CALLBACK,
+                  VOID_CB, page_size)
 from .entry import ArchiveEntry, new_archive_entry
 
 
@@ -43,6 +44,22 @@ def new_archive_read(format_name='all', filter_name='all'):
         ffi.read_free(archive_p)
 
 
+@contextmanager
+def custom_reader(
+        read_func, format_name='all', filter_name='all',
+        open_func=VOID_CB, close_func=VOID_CB, block_size=page_size,
+        archive_read_class=ArchiveRead
+):
+    """Read an archive using a custom function.
+    """
+    open_cb = OPEN_CALLBACK(open_func)
+    read_cb = READ_CALLBACK(read_func)
+    close_cb = CLOSE_CALLBACK(close_func)
+    with new_archive_read(format_name, filter_name) as archive_p:
+        ffi.read_open(archive_p, None, open_cb, read_cb, close_cb)
+        yield archive_read_class(archive_p)
+
+
 @contextmanager
 def fd_reader(fd, format_name='all', filter_name='all', block_size=4096):
     """Read an archive from a file descriptor.
@@ -76,3 +93,27 @@ def memory_reader(buf, format_name='all', filter_name='all'):
     with new_archive_read(format_name, filter_name) as archive_p:
         ffi.read_open_memory(archive_p, cast(buf, c_void_p), len(buf))
         yield ArchiveRead(archive_p)
+
+
+@contextmanager
+def stream_reader(stream, format_name='all', filter_name='all', block_size=page_size):
+    """Read an archive from a stream (an object supporting the `readinto` method).
+    """
+    buf = create_string_buffer(block_size)
+    buf_p = cast(buf, c_void_p)
+
+    def read_func(archive_p, context, ptrptr):
+        # readinto the buffer, returns number of bytes read
+        length = stream.readinto(buf)
+        # write the address of the buffer into the pointer
+        ptrptr = cast(ptrptr, POINTER(c_void_p))
+        ptrptr[0] = buf_p
+        # tell libarchive how much data was written into the buffer
+        return length
+
+    open_cb = OPEN_CALLBACK(VOID_CB)
+    read_cb = READ_CALLBACK(read_func)
+    close_cb = CLOSE_CALLBACK(VOID_CB)
+    with new_archive_read(format_name, filter_name) as archive_p:
+        ffi.read_open(archive_p, None, open_cb, read_cb, close_cb)
+        yield ArchiveRead(archive_p)

tests/test_rwx.py

@@ -1,6 +1,7 @@
 """Test reading, writing and extracting archives."""
 
 from __future__ import division, print_function, unicode_literals
+import io
 
 import libarchive
 from libarchive.extract import EXTRACT_OWNER, EXTRACT_PERM, EXTRACT_TIME
@@ -79,25 +80,18 @@ def test_files(tmpdir):
         assert tree2 == tree
 
 
-def test_custom_writer():
-
+def test_custom_writer_and_stream_reader():
     # Collect information on what should be in the archive
     tree = treestat('libarchive')
 
     # Create an archive of our libarchive/ directory
-    blocks = []
-
-    def write_cb(data):
-        blocks.append(data[:])
-        return len(data)
-
-    with libarchive.custom_writer(write_cb, 'zip') as archive:
+    stream = io.BytesIO()
+    with libarchive.custom_writer(stream.write, 'zip') as archive:
         archive.add_files('libarchive/')
-        pass
+    stream.seek(0)
 
     # Read the archive and check that the data is correct
-    buf = b''.join(blocks)
-    with libarchive.memory_reader(buf) as archive:
+    with libarchive.stream_reader(stream, 'zip') as archive:
         check_archive(archive, tree)