diff --git a/audeer/core/io.py b/audeer/core/io.py index 22a8eab..cc96f72 100644 --- a/audeer/core/io.py +++ b/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() diff --git a/tests/test_io_multithread.py b/tests/test_io_multithread.py new file mode 100644 index 0000000..39de8d1 --- /dev/null +++ b/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, + )