Extract archive into buffer (dict of io.BytesIO())

py7zr/compression.py

@@ -48,40 +48,43 @@ def __init__(self, files, src_start: int, header) -> None:
         self.files = files
         self.src_start = src_start
         self.header = header
+        self._dict = {}  # type: Dict[str, IO[Any]]
 
-    def extract(self, fp: BinaryIO, parallel: bool) -> None:
+    def extract(self, fp: BinaryIO, parallel: bool, return_dict: bool = False) -> None:
         """Extract worker method to handle 7zip folder and decompress each files."""
         if hasattr(self.header, 'main_streams') and self.header.main_streams is not None:
             src_end = self.src_start + self.header.main_streams.packinfo.packpositions[-1]
             numfolders = self.header.main_streams.unpackinfo.numfolders
             if numfolders == 1:
-                self.extract_single(fp, self.files, self.src_start, src_end)
+                self.extract_single(fp, self.files, self.src_start, src_end, return_dict)
             else:
                 folders = self.header.main_streams.unpackinfo.folders
                 positions = self.header.main_streams.packinfo.packpositions
                 empty_files = [f for f in self.files if f.emptystream]
                 if not parallel:
-                    self.extract_single(fp, empty_files, 0, 0)
+                    self.extract_single(fp, empty_files, 0, 0, return_dict)
                     for i in range(numfolders):
                         self.extract_single(fp, folders[i].files, self.src_start + positions[i],
-                                            self.src_start + positions[i + 1])
+                                            self.src_start + positions[i + 1], return_dict)
                 else:
                     filename = getattr(fp, 'name', None)
-                    self.extract_single(open(filename, 'rb'), empty_files, 0, 0)
+                    self.extract_single(open(filename, 'rb'), empty_files, 0, 0, return_dict)
                     extract_threads = []
                     for i in range(numfolders):
                         p = threading.Thread(target=self.extract_single,
                                              args=(filename, folders[i].files,
-                                                   self.src_start + positions[i], self.src_start + positions[i + 1]))
+                                                   self.src_start + positions[i], self.src_start + positions[i + 1],
+                                                   return_dict))
                         p.start()
                         extract_threads.append((p))
                     for p in extract_threads:
                         p.join()
         else:
             empty_files = [f for f in self.files if f.emptystream]
-            self.extract_single(fp, empty_files, 0, 0)
+            self.extract_single(fp, empty_files, 0, 0, return_dict)
 
-    def extract_single(self, fp: Union[BinaryIO, str], files, src_start: int, src_end: int) -> None:
+    def extract_single(self, fp: Union[BinaryIO, str], files, src_start: int, src_end: int,
+                       return_dict: bool = False) -> None:
         """Single thread extractor that takes file lists in single 7zip folder."""
         if files is None:
             return
@@ -91,12 +94,24 @@ def extract_single(self, fp: Union[BinaryIO, str], files, src_start: int, src_en
         for f in files:
             fileish = self.target_filepath.get(f.id, None)
             if fileish is not None:
-                with fileish.open(mode='wb') as ofp:
+                if return_dict:
+                    fname = str(fileish).replace("\\", "/")
+                    self._dict[fname] = io.BytesIO()
+                    ofp = self._dict[fname]
                     if not f.emptystream:
                         # extract to file
                         self.decompress(fp, f.folder, ofp, f.uncompressed[-1], f.compressed, src_end)
+                        ofp.seek(0)
                     else:
                         pass  # just create empty file
+                else:
+                    with fileish.open(mode='wb') as ofp:
+                        if not f.emptystream:
+                            # extract to file
+                            self.decompress(fp, f.folder, ofp, f.uncompressed[-1], f.compressed, src_end)
+                        else:
+                            pass  # just create empty file
+
             elif not f.emptystream:
                 # read and bin off a data but check crc
                 with NullIO() as ofp:

py7zr/py7zr.py

@@ -32,7 +32,7 @@
 import stat
 import sys
 from io import BytesIO
-from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union
+from typing import IO, Any, BinaryIO, Dict, List, Optional, Tuple, Union
 
 from py7zr.archiveinfo import Folder, Header, SignatureHeader
 from py7zr.compression import SevenZipCompressor, Worker, get_methods_names
@@ -644,15 +644,16 @@ def test(self) -> bool:
         """Test archive using CRC digests."""
         return self._test_digests()
 
-    def extractall(self, path: Optional[Any] = None) -> None:
+    def extractall(self, path: Optional[Any] = None, return_dict: bool = False) -> Optional[Dict[str, IO[Any]]]:
         """Extract all members from the archive to the current working
            directory and set owner, modification time and permissions on
            directories afterwards. `path' specifies a different directory
            to extract to.
         """
-        return self.extract(path)
+        return self.extract(path, return_dict=return_dict)
 
-    def extract(self, path: Optional[Any] = None, targets: Optional[List[str]] = None) -> None:
+    def extract(self, path: Optional[Any] = None, targets: Optional[List[str]] = None,
+                return_dict: bool = False) -> Optional[Dict[str, IO[Any]]]:
         target_junction = []  # type: List[pathlib.Path]
         target_sym = []  # type: List[pathlib.Path]
         target_files = []  # type: List[Tuple[pathlib.Path, Dict[str, Any]]]
@@ -724,7 +725,11 @@ def extract(self, path: Optional[Any] = None, targets: Optional[List[str]] = Non
                     raise Exception("Directory name is existed as a normal file.")
                 else:
                     raise Exception("Directory making fails on unknown condition.")
-        self.worker.extract(self.fp, parallel=(not self.password_protected and not self._filePassed))
+
+        self.worker.extract(self.fp, parallel=(not self.password_protected and not self._filePassed),
+                            return_dict=return_dict)
+        if return_dict:
+            return self.worker._dict
 
         # create symbolic links on target path as a working directory.
         # if path is None, work on current working directory.
@@ -746,6 +751,7 @@ def extract(self, path: Optional[Any] = None, targets: Optional[List[str]] = Non
 
         for o, p in target_files:
             self._set_file_property(o, p)
+        return None
 
     def writeall(self, path: Union[pathlib.Path, str], arcname: Optional[str] = None):
         """Write files in target path into archive."""

tests/__init__.py

@@ -34,7 +34,7 @@ def decode_all(archive, expected, tmpdir):
     for i, file_info in enumerate(archive.files):
         assert file_info.lastwritetime is not None
         assert file_info.filename is not None
-    archive.extractall(path=tmpdir)
+    archive.extractall(path=tmpdir, return_dict=False)
     archive.close()
     check_output(expected, tmpdir)