Improve the archive API with fileobj parameters

The archive API only allows unpacking an archive from file, and
writing an archive to file. However the zipfile, and tarfile
modules this API is based on, both accept fileobj parameters to
read from, and write to memory.

This change adds a fileobj parameter to create_archive, and
unpack_archive, in order to allow creating archives to, and
unpacking archives from memory.

Author: adanaja

src/e3/archive.py

@@ -10,7 +10,7 @@
 import tempfile
 import zipfile
 from contextlib import closing
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import e3
 import e3.error
@@ -20,7 +20,7 @@
 
 
 if TYPE_CHECKING:
-    from typing import Literal, Text, Union
+    from typing import Literal, Text, Union, IO, Any
     from collections.abc import Callable, Sequence
     from os import PathLike
     from e3.mypy import assert_never
@@ -146,6 +146,7 @@ def check_type(
 def unpack_archive(
     filename: str,
     dest: str,
+    fileobj: IO[bytes] | None = None,
     selected_files: Sequence[str] | None = None,
     remove_root_dir: RemoveRootDirType = False,
     unpack_cmd: Callable[..., None] | None = None,
@@ -159,6 +160,10 @@ def unpack_archive(
 
     :param filename: archive to unpack
     :param dest: destination directory (should exist)
+    :param fileobj: if specified, the archive is read from this file object
+        instead of opening a file. The file object must be opened in binary
+        mode. In this case filename is the name of the archive contained
+        in the file object.
     :param selected_files: list of files to unpack (partial extraction). If
         None all files are unpacked
     :param remove_root_dir: if True then the root dir of the archive is
@@ -192,7 +197,7 @@ def unpack_archive(
     logger.debug("unpack %s in %s", filename, dest)
     # First do some checks such as archive existence or destination directory
     # existence.
-    if not os.path.isfile(filename):
+    if fileobj is None and not os.path.isfile(filename):
         raise ArchiveError(origin="unpack_archive", message=f"cannot find {filename}")
 
     if not os.path.isdir(dest):
@@ -205,14 +210,19 @@ def unpack_archive(
 
     # We need to resolve to an absolute path as the extraction related
     # processes will be run in the destination directory
-    filename = os.path.abspath(filename)
+    if fileobj is None:
+        filename = os.path.abspath(filename)
 
     if unpack_cmd is not None:
         # Use user defined unpack command
-        if not selected_files:
-            return unpack_cmd(filename, dest)
-        else:
-            return unpack_cmd(filename, dest, selected_files=selected_files)
+        kwargs: dict[str, Any] = {}
+        if selected_files:
+            kwargs["selected_files"] = selected_files
+
+        if fileobj is not None:
+            kwargs["fileobj"] = fileobj
+
+        return unpack_cmd(filename, dest, **kwargs)
 
     ext = check_type(filename, force_extension=force_extension)
 
@@ -237,7 +247,13 @@ def unpack_archive(
                 elif ext.endswith("xz"):
                     mode += "xz"
                 # Extract tar files
-                with closing(tarfile.open(filename, mode=mode)) as fd:
+                with closing(
+                    tarfile.open(
+                        filename if fileobj is None else None,
+                        fileobj=fileobj,
+                        mode=mode,
+                    )
+                ) as fd:
                     check_selected = set(selected_files)
 
                     def is_match(name: str, files: Sequence[str]) -> bool:
@@ -291,7 +307,9 @@ def is_match(name: str, files: Sequence[str]) -> bool:
 
         elif ext == "zip":
             try:
-                with closing(E3ZipFile(filename, mode="r")) as zip_fd:
+                with closing(
+                    E3ZipFile(fileobj if fileobj is not None else filename, mode="r")
+                ) as zip_fd:
                     zip_fd.extractall(
                         tmp_dest, selected_files if selected_files else None
                     )
@@ -358,7 +376,8 @@ def is_match(name: str, files: Sequence[str]) -> bool:
 def create_archive(
     filename: str,
     from_dir: str,
-    dest: str,
+    dest: str | None = None,
+    fileobj: IO[bytes] | None = None,
     force_extension: str | None = None,
     from_dir_rename: str | None = None,
     no_root_dir: bool = False,
@@ -372,26 +391,45 @@ def create_archive(
 
     :param filename: archive to create
     :param from_dir: directory to pack (full path)
-    :param dest: destination directory (should exist)
+    :param dest: destination directory (should exist). If not specified,
+        the archive is written to the file object passed with fileobj.
+    :param fileobj: if specified, the archive is written to this file object
+        instead of opening a file. The file object must be opened in binary
+        mode. In this case filename is the name of the archive contained
+        in the file object.
     :param force_extension: specify the archive extension if not in the
         filename. If filename has no extension and force_extension is None
         create_archive will fail.
     :param from_dir_rename: name of root directory in the archive.
     :param no_root_dir: create archive without the root dir (zip only)
 
+    :raise ValueError: neither dest nor fileobj is provided
     :raise ArchiveError: if an error occurs
     """
+    if dest is None and fileobj is None:
+        raise ValueError("no destination provided")
+
     # Check extension
     from_dir = from_dir.rstrip("/")
-    filepath = os.path.abspath(os.path.join(dest, filename))
+
+    # If fileobj is None, dest is not None
+    filepath = (
+        os.path.abspath(os.path.join(cast(str, dest), filename))
+        if fileobj is None
+        else None
+    )
 
     ext = check_type(filename, force_extension=force_extension)
 
     if from_dir_rename is None:
         from_dir_rename = os.path.basename(from_dir)
 
     if ext == "zip":
-        zip_archive = zipfile.ZipFile(filepath, "w", zipfile.ZIP_DEFLATED)
+        zip_archive = zipfile.ZipFile(
+            cast(str, filepath) if fileobj is None else fileobj,
+            "w",
+            zipfile.ZIP_DEFLATED,
+        )
         for root, _, files in os.walk(from_dir):
             relative_root = os.path.relpath(
                 os.path.abspath(root), os.path.abspath(from_dir)
@@ -413,5 +451,7 @@ def create_archive(
             tar_format = "w:xz"
         else:
             assert_never()
-        with closing(tarfile.open(filepath, tar_format)) as tar_archive:
+        with closing(
+            tarfile.open(filepath, fileobj=fileobj, mode=tar_format)
+        ) as tar_archive:
             tar_archive.add(name=from_dir, arcname=from_dir_rename, recursive=True)

tests/tests_e3/archive/main_test.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import tempfile
+import io
 
 import e3.archive
 import e3.fs
@@ -91,6 +92,43 @@ def test_unpack(ext):
         e3.fs.rm(dest, True)
 
 
+@pytest.mark.parametrize("ext", (".tar.gz", ".zip"))
+def test_unpack_fileobj(ext):
+    dir_to_pack = os.path.dirname(__file__)
+
+    test_dir = os.path.basename(dir_to_pack)
+
+    dest = "dest"
+    e3.fs.mkdir(dest)
+
+    archive_name = "e3-core" + ext
+
+    try:
+        fo = io.BytesIO()
+        e3.archive.create_archive(
+            filename=archive_name,
+            from_dir=os.path.abspath(dir_to_pack),
+            fileobj=fo,
+        )
+
+        fo.seek(0)
+        e3.fs.mkdir(os.path.join(dest, "dest2"))
+        e3.archive.unpack_archive(
+            filename=archive_name,
+            dest=os.path.join(dest, "dest2"),
+            fileobj=fo,
+            selected_files=(
+                e3.os.fs.unixpath(os.path.join(test_dir, os.path.basename(__file__))),
+            ),
+            remove_root_dir=True,
+        )
+
+        assert os.path.exists(os.path.join(dest, "dest2", os.path.basename(__file__)))
+
+    finally:
+        e3.fs.rm(dest, True)
+
+
 def test_unsupported():
     """Test unsupported archive format."""
     with pytest.raises(e3.archive.ArchiveError) as err:
@@ -147,6 +185,48 @@ def custom_unpack(filename, dest, selected_files):
     assert t.kwargs["s"] == ["bar"]
 
 
+def test_unpack_cmd_fileobj():
+    """Test custom unpack_cmd with fileobj."""
+    dir_to_pack = os.path.dirname(__file__)
+
+    dest = "dest"
+    e3.fs.mkdir(dest)
+
+    archive_name = "e3-core.tar"
+
+    fo = io.BytesIO()
+    e3.archive.create_archive(
+        filename=archive_name,
+        from_dir=os.path.abspath(dir_to_pack),
+        fileobj=fo,
+    )
+
+    all_dest = "all_dest"
+    e3.fs.mkdir(all_dest)
+
+    # Use a custom unpack function and verify that it is called with
+    # the expected arguments
+    class TestResult:
+        def store_result(self, **kwargs):
+            self.kwargs = kwargs
+
+    t = TestResult()
+
+    def custom_unpack(filename, dest, fileobj):
+        t.store_result(f=filename, d=dest, fo=fileobj)
+
+    fo.seek(0)
+    e3.archive.unpack_archive(
+        filename=archive_name,
+        fileobj=fo,
+        dest=all_dest,
+        unpack_cmd=custom_unpack,
+    )
+    assert t.kwargs["f"] == archive_name
+    assert t.kwargs["d"] == all_dest
+    assert t.kwargs["fo"] == fo
+
+
 def test_unpack_files():
     """Test unpack_archive with selected_files."""
     e3.fs.mkdir("d")