diff --git a/audeer/core/io.py b/audeer/core/io.py index 22a8eab..b289c49 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 | None = 1, ) -> list[str]: r"""Extract ZIP or TAR file. @@ -302,6 +304,11 @@ 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 jobs or 1 for sequential + processing. If ``None`` will be set to the number of + processors on the machine multiplied by 5. + Multi-threading can significantly speed up extraction + for archives with many files Returns: paths of extracted files relative to ``destination`` @@ -365,24 +372,57 @@ def extract_archive( def extract_zip(archive: str) -> list: with zipfile.ZipFile(archive, "r") as zf: members = zf.infolist() - for member in progress_bar(members, desc=desc, disable=disable): - zf.extract(member, destination) + + if num_workers == 1: + # Sequential extraction with progress bar + for member in progress_bar(members, desc=desc, disable=disable): + zf.extract(member, destination) + else: + # Parallel extraction with progress bar + params = [([member, destination], {}) for member in members] + run_tasks( + zf.extract, + params, + num_workers=num_workers, + progress_bar=verbose, + task_description=desc, + ) + return [m.filename for m in members] + def extract_tar_member(member): + """Extract a single member from a TAR archive.""" + with tarfile.open(archive, "r") as tf: + # In Python 3.12 the `filter` argument was introduced, + # and it will be set automatically in Python 3.14, + # see + # https://docs.python.org/3.12/library/tarfile.html#tarfile-extraction-filter + # noqa: E501 + kwargs = {"numeric_owner": True} + if sys.version_info >= (3, 12): # pragma: no cover + kwargs = kwargs | {"filter": "tar"} + tf.extract(member, destination, **kwargs) + def extract_tar(archive: str) -> list: with tarfile.open(archive, "r") as tf: members = tf.getmembers() + + if num_workers == 1: + # Sequential extraction with progress bar 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, - # see - # https://docs.python.org/3.12/library/tarfile.html#tarfile-extraction-filter - # noqa: E501 - kwargs = {"numeric_owner": True} - if sys.version_info >= (3, 12): # pragma: no cover - kwargs = kwargs | {"filter": "tar"} - tf.extract(member, destination, **kwargs) - return [m.name for m in members] + extract_tar_member(member) + else: + # Parallel extraction with progress bar + params = [([member], {}) for member in members] + run_tasks( + extract_tar_member, + params, + num_workers=num_workers, + progress_bar=verbose, + task_description=desc, + ) + + return [m.name for m in members] try: if archive.endswith("zip"): diff --git a/tests/test_io.py b/tests/test_io.py index 0212e52..5213c56 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -1620,3 +1620,61 @@ def test_touch(tmpdir): new_stat = os.stat(path) assert stat.st_atime < new_stat.st_atime assert stat.st_mtime < new_stat.st_mtime + + +@pytest.mark.parametrize( + "archive_type", + [ + "zip", + "tar.gz", + ], +) +def test_extract_archive_multithread(tmpdir, archive_type): + """Test multi-threaded extraction of archives.""" + # Create test files + root = audeer.mkdir(tmpdir, "root") + files = [] + for i in range(20): + file = audeer.touch(root, f"file{i:02d}.txt") + with open(file, "w") as f: + f.write(f"Content of file {i}") + files.append(f"file{i:02d}.txt") + + # Create archive + archive = audeer.path(tmpdir, f"archive.{archive_type}") + audeer.create_archive(root, None, archive) + + # Extract with single thread (default) + dest_single = audeer.mkdir(tmpdir, "dest_single") + result_single = audeer.extract_archive( + archive, + dest_single, + num_workers=1, + ) + + # Extract with multiple threads + dest_multi = audeer.mkdir(tmpdir, "dest_multi") + result_multi = audeer.extract_archive( + archive, + dest_multi, + num_workers=4, + ) + + # Verify results are identical + assert result_single == result_multi + assert len(result_single) == 20 + + # Verify all files were extracted correctly in both cases + for file in files: + single_file = audeer.path(dest_single, file) + multi_file = audeer.path(dest_multi, file) + + assert os.path.exists(single_file) + assert os.path.exists(multi_file) + + with open(single_file) as f: + content_single = f.read() + with open(multi_file) as f: + content_multi = f.read() + + assert content_single == content_multi