Multithreading archive extraction

audeer/core/io.py

@@ -15,6 +15,7 @@
 from audeer.core.path import path as safe_path
 from audeer.core.tqdm import format_display_message
 from audeer.core.tqdm import progress_bar
+from audeer.core.utils import run_tasks
 from audeer.core.utils import to_list
 
 
@@ -292,6 +293,7 @@ def extract_archive(
     *,
     keep_archive: bool = True,
     verbose: bool = False,
+    num_workers: int = 1,
 ) -> list[str]:
     r"""Extract ZIP or TAR file.
 
@@ -302,6 +304,8 @@ def extract_archive(
             it will be created
         keep_archive: if ``False`` delete archive file after extraction
         verbose: if ``True`` a progress bar is shown
+        num_workers: number of parallel workers for extraction.
+            Must be at least 1
 
     Returns:
         paths of extracted files relative to ``destination``
@@ -313,6 +317,7 @@ def extract_archive(
         NotADirectoryError: if ``destination`` is not a directory
         RuntimeError: if ``archive`` is not a ZIP or TAR file
         RuntimeError: if ``archive`` is malformed
+        ValueError: if ``num_workers`` is less than 1
 
     Examples:
         >>> folder = audeer.path(".")
@@ -328,6 +333,9 @@ def extract_archive(
     archive = safe_path(archive)
     destination = safe_path(destination)
 
+    if num_workers < 1:
+        raise ValueError(f"num_workers must be at least 1, got {num_workers}")
+
     if not os.path.exists(archive):
         raise FileNotFoundError(
             errno.ENOENT,
@@ -362,16 +370,32 @@ def extract_archive(
     )
     disable = not verbose
 
-    def extract_zip(archive: str) -> list:
+    def extract_zip(
+        archive: str,
+        members: Sequence[str | zipfile.ZipInfo] | None = None,
+    ) -> list:
         with zipfile.ZipFile(archive, "r") as zf:
-            members = zf.infolist()
+            if members is None:
+                members = zf.infolist()
+            else:
+                members = [zf.getinfo(m) for m in members]
             for member in progress_bar(members, desc=desc, disable=disable):
                 zf.extract(member, destination)
             return [m.filename for m in members]
 
-    def extract_tar(archive: str) -> list:
+    def extract_tar(
+        archive: str,
+        members: Sequence[str | tarfile.TarInfo] | None = None,
+    ) -> list:
         with tarfile.open(archive, "r") as tf:
-            members = tf.getmembers()
+            if members is None:
+                members = tf.getmembers()
+            else:
+                resolved_members = []
+                for m in members:
+                    member = tf.getmember(m)
+                    resolved_members.append(member)
+                members = resolved_members
             for member in progress_bar(members, desc=desc, disable=disable):
                 # In Python 3.12 the `filter` argument was introduced,
                 # and it will be set automatically in Python 3.14,
@@ -384,15 +408,75 @@ def extract_tar(archive: str) -> list:
                 tf.extract(member, destination, **kwargs)
             return [m.name for m in members]
 
+    def extract_chunk(
+        archive: str,
+        members: Sequence[str],
+        destination: str,
+        extension: str,
+    ) -> list[str]:
+        if extension == "zip":
+            return extract_zip(archive, members)
+        else:
+            return extract_tar(archive, members)
+
     try:
         if archive.endswith("zip"):
-            files = extract_zip(archive)
+            extension = "zip"
+            if num_workers > 1:
+                with zipfile.ZipFile(archive, "r") as zf:
+                    members = zf.namelist()
+            else:
+                files = extract_zip(archive)
         elif tarfile.is_tarfile(archive):
-            files = extract_tar(archive)
+            extension = "tar"
+            if num_workers > 1:
+                with tarfile.open(archive, "r") as tf:
+                    members = tf.getnames()
+            else:
+                files = extract_tar(archive)
         else:
             raise RuntimeError(f"You can only extract ZIP and TAR files, not {archive}")
+
+        if num_workers > 1:
+            # Handle empty archives
+            if not members:
+                files = []
+            else:
+                # Use round-robin distribution for better load balancing
+                # Track original indices to preserve order after extraction
+                chunks = [[] for _ in range(num_workers)]
+                indices = [[] for _ in range(num_workers)]
+                for i, member in enumerate(members):
+                    worker_idx = i % num_workers
+                    chunks[worker_idx].append(member)
+                    indices[worker_idx].append(i)
+
+                # Filter out empty chunks (when num_workers > len(members))
+                non_empty = [(c, idx) for c, idx in zip(chunks, indices) if c]
+                chunks = [c for c, _ in non_empty]
+                indices = [idx for _, idx in non_empty]
+
+                params = [
+                    ([archive, chunk, destination, extension], {}) for chunk in chunks
+                ]
+                chunk_results = run_tasks(
+                    extract_chunk,
+                    params,
+                    num_workers=len(chunks),
+                    progress_bar=verbose,
+                    task_description=f"Extract {os.path.basename(archive)}",
+                )
+
+                # Reconstruct files in original order
+                files = [None] * len(members)
+                for chunk_indices, chunk_files in zip(indices, chunk_results):
+                    for idx, file in zip(chunk_indices, chunk_files):
+                        files[idx] = file
+
     except (EOFError, zipfile.BadZipFile, tarfile.ReadError):
         raise RuntimeError(f"Broken archive: {archive}")
+    except KeyError as e:
+        raise RuntimeError(f"Member not found in archive {archive}: {e}")
     except (KeyboardInterrupt, Exception):  # pragma: no cover
         # Clean up broken extraction files
         if destination_created:
@@ -416,6 +500,7 @@ def extract_archives(
     *,
     keep_archive: bool = True,
     verbose: bool = False,
+    num_workers: int = 1,
 ) -> list[str]:
     r"""Extract multiple ZIP or TAR archives at once.
 
@@ -426,6 +511,8 @@ def extract_archives(
             it will be created
         keep_archive: if ``False`` delete archive files after extraction
         verbose: if ``True`` a progress bar is shown
+        num_workers: number of parallel workers for extraction.
+            Must be at least 1
 
     Returns:
         paths of extracted files relative to ``destination``
@@ -437,6 +524,7 @@ def extract_archives(
         NotADirectoryError: if ``destination`` is not a directory
         RuntimeError: if an archive is not a ZIP or TAR file
         RuntimeError: if an archive file is malformed
+        ValueError: if ``num_workers`` is less than 1
 
     Examples:
         >>> folder = audeer.path(".")
@@ -467,6 +555,7 @@ def extract_archives(
                 destination,
                 keep_archive=keep_archive,
                 verbose=False,
+                num_workers=num_workers,
             )
             pbar.update()
 

tests/test_io_multithread.py

@@ -0,0 +1,191 @@
+import os
+from unittest import mock
+import zipfile
+
+import pytest
+
+import audeer
+
+
+@pytest.fixture
+def archive_files():
+    """List of files to create in test archives."""
+    return [f"file_{i}.txt" for i in range(10)]
+
+
+@pytest.fixture
+def zip_archive(tmpdir, archive_files):
+    """Create a zip archive with test files."""
+    root = audeer.mkdir(tmpdir, "root_zip")
+    for file in archive_files:
+        audeer.touch(root, file)
+    archive = audeer.path(tmpdir, "archive.zip")
+    audeer.create_archive(root, archive_files, archive)
+    return archive, archive_files
+
+
+@pytest.fixture
+def tar_archive(tmpdir, archive_files):
+    """Create a tar archive with test files."""
+    root = audeer.mkdir(tmpdir, "root_tar")
+    for file in archive_files:
+        audeer.touch(root, file)
+    archive = audeer.path(tmpdir, "archive.tar")
+    audeer.create_archive(root, archive_files, archive)
+    return archive, archive_files
+
+
+@pytest.fixture
+def empty_zip_archive(tmpdir):
+    """Create an empty zip archive."""
+    root = audeer.mkdir(tmpdir, "root_empty_zip")
+    archive = audeer.path(tmpdir, "empty.zip")
+    audeer.create_archive(root, [], archive)
+    return archive
+
+
+@pytest.fixture
+def empty_tar_archive(tmpdir):
+    """Create an empty tar archive."""
+    root = audeer.mkdir(tmpdir, "root_empty_tar")
+    archive = audeer.path(tmpdir, "empty.tar")
+    audeer.create_archive(root, [], archive)
+    return archive
+
+
+@pytest.mark.parametrize("num_workers", [1, 2, 4])
+def test_extract_zip_multithread(tmpdir, zip_archive, num_workers):
+    archive, expected_files = zip_archive
+    destination = audeer.mkdir(tmpdir, f"extracted_zip_{num_workers}")
+
+    extracted_files = audeer.extract_archive(
+        archive,
+        destination,
+        num_workers=num_workers,
+    )
+
+    assert len(extracted_files) == len(expected_files)
+    assert extracted_files == expected_files
+    extracted_paths = [os.path.join(destination, f) for f in expected_files]
+    assert all(os.path.exists(p) for p in extracted_paths)
+
+
+@pytest.mark.parametrize("num_workers", [1, 2, 4])
+def test_extract_tar_multithread(tmpdir, tar_archive, num_workers):
+    archive, expected_files = tar_archive
+    destination = audeer.mkdir(tmpdir, f"extracted_tar_{num_workers}")
+
+    extracted_files = audeer.extract_archive(
+        archive,
+        destination,
+        num_workers=num_workers,
+    )
+
+    assert len(extracted_files) == len(expected_files)
+    assert extracted_files == expected_files
+    extracted_paths = [os.path.join(destination, f) for f in expected_files]
+    assert all(os.path.exists(p) for p in extracted_paths)
+
+
+@pytest.mark.parametrize("num_workers", [1, 2])
+def test_extract_empty_zip_archive(tmpdir, empty_zip_archive, num_workers):
+    destination = audeer.mkdir(tmpdir, f"extracted_empty_zip_{num_workers}")
+
+    extracted_files = audeer.extract_archive(
+        empty_zip_archive,
+        destination,
+        num_workers=num_workers,
+    )
+
+    assert extracted_files == []
+
+
+@pytest.mark.parametrize("num_workers", [1, 2])
+def test_extract_empty_tar_archive(tmpdir, empty_tar_archive, num_workers):
+    destination = audeer.mkdir(tmpdir, f"extracted_empty_tar_{num_workers}")
+
+    extracted_files = audeer.extract_archive(
+        empty_tar_archive,
+        destination,
+        num_workers=num_workers,
+    )
+
+    assert extracted_files == []
+
+
+def test_extract_archives_multithread(tmpdir, archive_files):
+    root = audeer.mkdir(tmpdir, "root_multi")
+    for file in archive_files:
+        audeer.touch(root, file)
+
+    archive1 = audeer.path(tmpdir, "archive1.zip")
+    audeer.create_archive(root, archive_files, archive1)
+    archive2 = audeer.path(tmpdir, "archive2.zip")
+    audeer.create_archive(root, archive_files, archive2)
+
+    destination = audeer.mkdir(tmpdir, "extracted_archives")
+
+    extracted_files = audeer.extract_archives(
+        [archive1, archive2],
+        destination,
+        num_workers=2,
+    )
+
+    # extract_archives returns a flat list of all extracted files
+    # Since we extract to the same destination and files have same names,
+    # they overwrite. But the returned list should contain all of them.
+    assert len(extracted_files) == 20
+    assert len(set(extracted_files)) == 10  # duplicates because same filenames
+
+
+@pytest.mark.parametrize("num_workers", [0, -1, -10])
+def test_extract_archive_invalid_num_workers(tmpdir, zip_archive, num_workers):
+    archive, _ = zip_archive
+    destination = audeer.mkdir(tmpdir, f"extracted_invalid_{num_workers}")
+
+    with pytest.raises(ValueError, match="num_workers must be at least 1"):
+        audeer.extract_archive(
+            archive,
+            destination,
+            num_workers=num_workers,
+        )
+
+
+def test_extract_archive_preserves_order(tmpdir):
+    """Test that extraction preserves original file order."""
+    root = audeer.mkdir(tmpdir, "root_order")
+    # Create files with names that would sort differently alphabetically
+    files = ["z_first.txt", "a_second.txt", "m_third.txt"]
+    for file in files:
+        audeer.touch(root, file)
+
+    archive = audeer.path(tmpdir, "order_test.zip")
+    audeer.create_archive(root, files, archive)
+
+    destination = audeer.mkdir(tmpdir, "extracted_order")
+
+    extracted_files = audeer.extract_archive(
+        archive,
+        destination,
+        num_workers=2,
+    )
+
+    # Files should be returned in the order they were added to the archive
+    assert extracted_files == files
+
+
+def test_extract_archive_member_not_found(tmpdir, zip_archive):
+    """Test that KeyError from missing member is converted to RuntimeError."""
+    archive, _ = zip_archive
+    destination = audeer.mkdir(tmpdir, "extracted_missing")
+
+    def mock_getinfo(self, name):
+        raise KeyError(name)
+
+    with mock.patch.object(zipfile.ZipFile, "getinfo", mock_getinfo):
+        with pytest.raises(RuntimeError, match="Member not found in archive"):
+            audeer.extract_archive(
+                archive,
+                destination,
+                num_workers=2,
+            )