From 20c9a2e4d9e24ae9c04ce4a2f547b815823dcad5 Mon Sep 17 00:00:00 2001 From: Hagen Wierstorf Date: Wed, 19 Nov 2025 11:26:13 +0100 Subject: [PATCH 1/2] Multithreading archive extraction --- audeer/core/io.py | 85 +++++++++++++++++++++++++++++++++--- tests/test_io_multithread.py | 78 +++++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 6 deletions(-) create mode 100644 tests/test_io_multithread.py diff --git a/audeer/core/io.py b/audeer/core/io.py index 22a8eab..6283f8b 100644 --- a/audeer/core/io.py +++ b/audeer/core/io.py @@ -15,6 +15,8 @@ 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 flatten_list +from audeer.core.utils import run_tasks from audeer.core.utils import to_list @@ -292,6 +294,7 @@ def extract_archive( *, keep_archive: bool = True, verbose: bool = False, + num_workers: int = 1, ) -> list[str]: r"""Extract ZIP or TAR file. @@ -301,7 +304,15 @@ def extract_archive( If the folder does not exists, it will be created keep_archive: if ``False`` delete archive file after extraction + keep_archive: if ``False`` delete archive file after extraction verbose: if ``True`` a progress bar is shown + num_workers: number of workers to use for extraction. + If ``num_workers > 1``, + the archive is split into chunks + and extracted in parallel. + This is only supported for ZIP and TAR files. + If ``num_workers = 1``, + the archive is extracted sequentially Returns: paths of extracted files relative to ``destination`` @@ -362,16 +373,28 @@ def extract_archive( ) disable = not verbose - def extract_zip(archive: str) -> list: + def extract_zip( + archive: str, + members: Sequence[str] | 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] | None = None, + ) -> list: with tarfile.open(archive, "r") as tf: - members = tf.getmembers() + if members is None: + members = tf.getmembers() + else: + members = [tf.getmember(m) for m in 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,13 +407,53 @@ 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: + # Split members into chunks + chunk_size = len(members) // num_workers + 1 + chunks = [ + members[i : i + chunk_size] for i in range(0, len(members), chunk_size) + ] + params = [ + ([archive, chunk, destination, extension], {}) for chunk in chunks + ] + files = run_tasks( + extract_chunk, + params, + num_workers=num_workers, + progress_bar=verbose, + task_description=f"Extract {os.path.basename(archive)}", + ) + files = flatten_list(files) + except (EOFError, zipfile.BadZipFile, tarfile.ReadError): raise RuntimeError(f"Broken archive: {archive}") except (KeyboardInterrupt, Exception): # pragma: no cover @@ -416,6 +479,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. @@ -425,7 +489,15 @@ def extract_archives( If the folder does not exists, it will be created keep_archive: if ``False`` delete archive files after extraction + keep_archive: if ``False`` delete archive files after extraction verbose: if ``True`` a progress bar is shown + num_workers: number of workers to use for extraction. + If ``num_workers > 1``, + the archive is split into chunks + and extracted in parallel. + This is only supported for ZIP and TAR files. + If ``num_workers = 1``, + the archive is extracted sequentially Returns: paths of extracted files relative to ``destination`` @@ -467,6 +539,7 @@ def extract_archives( destination, keep_archive=keep_archive, verbose=False, + num_workers=num_workers, ) pbar.update() diff --git a/tests/test_io_multithread.py b/tests/test_io_multithread.py new file mode 100644 index 0000000..772b337 --- /dev/null +++ b/tests/test_io_multithread.py @@ -0,0 +1,78 @@ +import os + +import audeer + + +def test_extract_archive_multithread(tmpdir): + root = audeer.mkdir(tmpdir, "root") + # Create a bunch of files + files = [f"file_{i}.txt" for i in range(10)] + for file in files: + audeer.touch(root, file) + + archive = audeer.path(tmpdir, "archive.zip") + audeer.create_archive(root, files, archive) + + destination = audeer.mkdir(tmpdir, "extracted") + + # Extract with num_workers=2 + extracted_files = audeer.extract_archive( + archive, + destination, + num_workers=2, + ) + + assert len(extracted_files) == 10 + assert sorted(extracted_files) == sorted(files) + for file in files: + assert os.path.exists(os.path.join(destination, file)) + + +def test_extract_tar_multithread(tmpdir): + root = audeer.mkdir(tmpdir, "root") + files = [f"file_{i}.txt" for i in range(10)] + for file in files: + audeer.touch(root, file) + + archive = audeer.path(tmpdir, "archive.tar") + audeer.create_archive(root, files, archive) + + destination = audeer.mkdir(tmpdir, "extracted_tar") + + extracted_files = audeer.extract_archive( + archive, + destination, + num_workers=2, + ) + + assert len(extracted_files) == 10 + assert sorted(extracted_files) == sorted(files) + for file in files: + assert os.path.exists(os.path.join(destination, file)) + + +def test_extract_archives_multithread(tmpdir): + root = audeer.mkdir(tmpdir, "root") + files = [f"file_{i}.txt" for i in range(10)] + for file in files: + audeer.touch(root, file) + + archive1 = audeer.path(tmpdir, "archive1.zip") + audeer.create_archive(root, files, archive1) + archive2 = audeer.path(tmpdir, "archive2.zip") + audeer.create_archive(root, 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 From 7c1782610a8be775c7b981de32a8c28d65eb745d Mon Sep 17 00:00:00 2001 From: Hagen Wierstorf Date: Thu, 8 Jan 2026 09:19:16 +0100 Subject: [PATCH 2/2] Improve code --- audeer/core/io.py | 88 ++++++++++------- tests/test_io_multithread.py | 185 ++++++++++++++++++++++++++++------- 2 files changed, 201 insertions(+), 72 deletions(-) diff --git a/audeer/core/io.py b/audeer/core/io.py index 6283f8b..cc96f72 100644 --- a/audeer/core/io.py +++ b/audeer/core/io.py @@ -15,7 +15,6 @@ 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 flatten_list from audeer.core.utils import run_tasks from audeer.core.utils import to_list @@ -304,15 +303,9 @@ def extract_archive( If the folder does not exists, it will be created keep_archive: if ``False`` delete archive file after extraction - keep_archive: if ``False`` delete archive file after extraction verbose: if ``True`` a progress bar is shown - num_workers: number of workers to use for extraction. - If ``num_workers > 1``, - the archive is split into chunks - and extracted in parallel. - This is only supported for ZIP and TAR files. - If ``num_workers = 1``, - the archive is extracted sequentially + num_workers: number of parallel workers for extraction. + Must be at least 1 Returns: paths of extracted files relative to ``destination`` @@ -324,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(".") @@ -339,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, @@ -375,7 +372,7 @@ def extract_archive( def extract_zip( archive: str, - members: Sequence[str] | None = None, + members: Sequence[str | zipfile.ZipInfo] | None = None, ) -> list: with zipfile.ZipFile(archive, "r") as zf: if members is None: @@ -388,13 +385,17 @@ def extract_zip( def extract_tar( archive: str, - members: Sequence[str] | None = None, + members: Sequence[str | tarfile.TarInfo] | None = None, ) -> list: with tarfile.open(archive, "r") as tf: if members is None: members = tf.getmembers() else: - members = [tf.getmember(m) for m in members] + 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, @@ -437,25 +438,45 @@ def extract_chunk( raise RuntimeError(f"You can only extract ZIP and TAR files, not {archive}") if num_workers > 1: - # Split members into chunks - chunk_size = len(members) // num_workers + 1 - chunks = [ - members[i : i + chunk_size] for i in range(0, len(members), chunk_size) - ] - params = [ - ([archive, chunk, destination, extension], {}) for chunk in chunks - ] - files = run_tasks( - extract_chunk, - params, - num_workers=num_workers, - progress_bar=verbose, - task_description=f"Extract {os.path.basename(archive)}", - ) - files = flatten_list(files) + # 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: @@ -489,15 +510,9 @@ def extract_archives( If the folder does not exists, it will be created keep_archive: if ``False`` delete archive files after extraction - keep_archive: if ``False`` delete archive files after extraction verbose: if ``True`` a progress bar is shown - num_workers: number of workers to use for extraction. - If ``num_workers > 1``, - the archive is split into chunks - and extracted in parallel. - This is only supported for ZIP and TAR files. - If ``num_workers = 1``, - the archive is extracted sequentially + num_workers: number of parallel workers for extraction. + Must be at least 1 Returns: paths of extracted files relative to ``destination`` @@ -509,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(".") diff --git a/tests/test_io_multithread.py b/tests/test_io_multithread.py index 772b337..39de8d1 100644 --- a/tests/test_io_multithread.py +++ b/tests/test_io_multithread.py @@ -1,66 +1,127 @@ import os +from unittest import mock +import zipfile + +import pytest import audeer -def test_extract_archive_multithread(tmpdir): - root = audeer.mkdir(tmpdir, "root") - # Create a bunch of files - files = [f"file_{i}.txt" for i in range(10)] - for file in files: - audeer.touch(root, file) +@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, files, archive) + 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 + - destination = audeer.mkdir(tmpdir, "extracted") +@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}") - # Extract with num_workers=2 extracted_files = audeer.extract_archive( archive, destination, - num_workers=2, + num_workers=num_workers, ) - assert len(extracted_files) == 10 - assert sorted(extracted_files) == sorted(files) - for file in files: - assert os.path.exists(os.path.join(destination, file)) + 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) -def test_extract_tar_multithread(tmpdir): - root = audeer.mkdir(tmpdir, "root") - files = [f"file_{i}.txt" for i in range(10)] - for file in files: - audeer.touch(root, file) +@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) - archive = audeer.path(tmpdir, "archive.tar") - audeer.create_archive(root, files, archive) - destination = audeer.mkdir(tmpdir, "extracted_tar") +@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( - archive, + empty_zip_archive, destination, - num_workers=2, + num_workers=num_workers, ) - assert len(extracted_files) == 10 - assert sorted(extracted_files) == sorted(files) - for file in files: - assert os.path.exists(os.path.join(destination, file)) + assert extracted_files == [] -def test_extract_archives_multithread(tmpdir): - root = audeer.mkdir(tmpdir, "root") - files = [f"file_{i}.txt" for i in range(10)] - for file in 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, files, archive1) + audeer.create_archive(root, archive_files, archive1) archive2 = audeer.path(tmpdir, "archive2.zip") - audeer.create_archive(root, files, archive2) + audeer.create_archive(root, archive_files, archive2) destination = audeer.mkdir(tmpdir, "extracted_archives") @@ -72,7 +133,59 @@ def test_extract_archives_multithread(tmpdir): # 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. + # 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, + )