diff --git a/docs/index.md b/docs/index.md index 62819fd..9818e3e 100644 --- a/docs/index.md +++ b/docs/index.md @@ -102,12 +102,21 @@ Use `download_dataset` to download a dataset to the configured download path: ```python from datacollective import download_dataset -dataset = download_dataset("your-dataset-id") +archive_path = download_dataset("your-dataset-id") +``` + +The function returns the path to the downloaded archive. Pass `extract=True` to also unpack +it into a directory beside the archive, named after the archive without its extension, and +get that directory back instead: -# Depending on the implementation, `dataset` may contain metadata -# about the downloaded files or a higher-level dataset object. +```python +dataset_dir = download_dataset("your-dataset-id", extract=True) ``` +An existing archive is reused unless `overwrite_existing=True`, and an existing extracted +directory is reused unless `overwrite_extracted=True` (or the archive was downloaded again). +This is the same directory `load_dataset` extracts to, so the two never unpack an archive twice. + !!! note `download_dataset` was previously called `save_dataset_to_disk`. The old name still works for backward compatibility, but it is deprecated and new code should use `download_dataset`. @@ -128,11 +137,17 @@ The sample archive is saved under `samples//` inside `MDC_DOWNLOAD_P apart from the full dataset archives. A sample may share a filename with the full archive, so keeping them in separate directories ensures one is never mistaken for the other. -The function returns the path to the sample archive. The archive is not extracted; use -`load_dataset` for the full dataset if you want a loaded DataFrame. +The function returns the path to the sample archive. Pass `extract=True` to also unpack it +into a directory beside the archive and get that directory back instead: + +```python +sample_dir = download_sample("your-dataset-id", extract=True) +``` If the sample is already present, it is returned without downloading again. Pass -`overwrite_existing=True` to download it again and replace the existing file. +`overwrite_existing=True` to download it again and replace the existing file. With +`extract=True`, an existing extracted directory is likewise reused unless +`overwrite_extracted=True` or the sample was downloaded again. !!! note Downloading a sample requires an API key whose organization has agreed to the dataset's diff --git a/src/datacollective/archive_utils.py b/src/datacollective/archive_utils.py index 5521baa..51722d1 100644 --- a/src/datacollective/archive_utils.py +++ b/src/datacollective/archive_utils.py @@ -17,6 +17,12 @@ def _extract_archive( Extract the given `.tar.gz` / `.tgz` archive into `dest_dir`. If the extracted directory already exists and overwrite_extracted is False, skip extraction. + The archive is extracted under a temporary name and renamed once complete, so an + interrupted extraction is never mistaken for a finished one. When overwriting, the + existing directory is replaced only after the new extraction has succeeded, so a bad + archive never costs the user a good directory. A failed extraction cleans up after + itself. + Args: archive_path: Path to the archive file. dest_dir: Directory where to extract the contents. @@ -26,27 +32,57 @@ def _extract_archive( Raises: ValueError: If the archive is not a `.tar.gz` / `.tgz` file. + FileExistsError: If something other than a directory sits where the archive + would be extracted. It is never deleted, even with `overwrite_extracted`. """ - suffix = next((s for s in TAR_GZ_SUFFIXES if archive_path.name.endswith(s)), None) - if suffix is None: + dir_name = _extracted_dir_name(archive_path.name) + if dir_name is None: raise ValueError( f"Unsupported archive type for `{archive_path.name}`. " f"Expected {' or '.join(TAR_GZ_SUFFIXES)}." ) # Extract into a dedicated directory under `dest_dir` named after the archive - target = dest_dir / archive_path.name.removesuffix(suffix) - if target.exists(): - if not overwrite_extracted: - logger.info( - f"Extracted directory already exists. Skipping extraction: `{target}`" + target = dest_dir / dir_name + partial = target.with_name(target.name + ".partial") + # Only directories this function created may be reused or removed. A stray file in + # the way cannot be the extracted dataset, and deleting it would destroy user data. + for path in (target, partial): + if path.exists() and not path.is_dir(): + raise FileExistsError( + f"Cannot extract `{archive_path.name}`: `{path}` exists and is not a " + "directory. Move or remove it first." ) - return target - logger.info(f"Overwriting existing extracted directory: `{target}`") - shutil.rmtree(target) + if target.exists() and not overwrite_extracted: + logger.info( + f"Extracted directory already exists. Skipping extraction: `{target}`" + ) + return target + + # Left over from an interrupted extraction; start again. + if partial.exists(): + shutil.rmtree(partial) - target.mkdir(parents=True, exist_ok=True) - with tarfile.open(archive_path, "r:gz") as tf: - tf.extractall(path=target, filter="data") + partial.mkdir(parents=True) + try: + with tarfile.open(archive_path, "r:gz") as tf: + tf.extractall(path=partial, filter="data") + # Only now, with a complete extraction in hand, is the old directory given up. + if target.exists(): + logger.info(f"Overwriting existing extracted directory: `{target}`") + shutil.rmtree(target) + partial.rename(target) + except BaseException: + shutil.rmtree(partial, ignore_errors=True) + raise return target + + +def _extracted_dir_name(archive_name: str) -> str | None: + """Name of the directory an archive extracts into: its name without the archive + suffix. None if it is not a supported archive.""" + suffix = next((s for s in TAR_GZ_SUFFIXES if archive_name.endswith(s)), None) + if suffix is None or archive_name == suffix: + return None + return archive_name.removesuffix(suffix) diff --git a/src/datacollective/datasets.py b/src/datacollective/datasets.py index e1dd2a9..eeedc40 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -85,6 +85,8 @@ def download_dataset( show_progress: bool = True, overwrite_existing: bool = False, enable_logging: bool = False, + extract: bool = False, + overwrite_extracted: bool = False, ) -> Path: """ Download the dataset archive to a local directory and return the archive path. @@ -93,6 +95,12 @@ def download_dataset( Automatically resumes interrupted downloads if a matching .checksum file exists from a previous attempt. + With `extract=True`, also unpacks the archive into a directory beside it, named after + the archive without its extension, and returns that directory instead. An existing + directory is reused unless `overwrite_extracted=True`, or the archive itself was + downloaded again with `overwrite_existing=True`. This is the directory `load_dataset` + extracts to, so the two never unpack the same archive twice. + Note: Previously called `save_dataset_to_disk`, which remains available as a deprecated alias for backward compatibility. @@ -101,14 +109,22 @@ def download_dataset( download_directory: Directory where to save the downloaded archive file. If None or empty, falls back to env MDC_DOWNLOAD_PATH or default. show_progress: Whether to show a progress bar during download. - overwrite_existing: Whether to overwrite the existing archive file. + overwrite_existing: Whether to overwrite the existing archive file. With + `extract=True` the extracted directory is then unpacked afresh too. enable_logging: Whether to enable SDK logging to console and a local log file. + extract: Whether to also unpack the archive and return the extracted directory. + overwrite_extracted: With `extract=True`, whether to unpack the archive again even + if the extracted directory already exists. Does not download the archive again. Returns: - Path to the downloaded dataset archive. + Path to the downloaded dataset archive, or to the extracted directory if + `extract=True`. Raises: - ValueError: If dataset_id is empty. + ValueError: If dataset_id is empty, or `extract=True` and the archive is not a + `.tar.gz` / `.tgz` file. + FileExistsError: If `extract=True` and a file, rather than a directory, sits + where the archive would be extracted. FileNotFoundError: If the dataset does not exist (404). ResourceRemovedError: If the dataset is no longer available (410). PermissionError: If access is denied (403) or download directory is not writable. @@ -122,6 +138,8 @@ def download_dataset( overwrite_existing, enable_logging, download_source=DOWNLOAD_SOURCE_DOWNLOAD, + extract=extract, + overwrite_extracted=overwrite_extracted, ) @@ -131,6 +149,8 @@ def download_sample( show_progress: bool = True, overwrite_existing: bool = False, enable_logging: bool = False, + extract: bool = False, + overwrite_extracted: bool = False, ) -> Path: """ Download the dataset's sample archive to a local directory and return the archive path. @@ -145,6 +165,11 @@ def download_sample( Unlike `download_dataset`, an interrupted sample download starts again rather than resuming: the API returns no checksum for samples to match a partial file against. + With `extract=True`, also unpacks the archive into a directory beside it, named after + the archive without its extension, and returns that directory instead. An existing + directory is reused unless `overwrite_extracted=True`, or the sample itself was + downloaded again with `overwrite_existing=True`. + Requires an API key whose organization has agreed to the dataset's terms on the MDC platform. Buying the dataset is not required. @@ -154,13 +179,21 @@ def download_sample( If None or empty, falls back to env MDC_DOWNLOAD_PATH or default. show_progress: Whether to show a progress bar during download. overwrite_existing: Whether to download the sample again, replacing the existing one. + With `extract=True` the extracted directory is then unpacked afresh too. enable_logging: Whether to enable SDK logging to console and a local log file. + extract: Whether to also unpack the archive and return the extracted directory. + overwrite_extracted: With `extract=True`, whether to unpack the sample again even + if the extracted directory already exists. Does not download the sample again. Returns: - Path to the downloaded sample archive. + Path to the downloaded sample archive, or to the extracted directory if + `extract=True`. Raises: - ValueError: If dataset_id is empty. + ValueError: If dataset_id is empty, or `extract=True` and the sample is not a + `.tar.gz` / `.tgz` file. + FileExistsError: If `extract=True` and a file, rather than a directory, sits + where the sample would be extracted. FileNotFoundError: If the dataset does not exist, or publishes no sample (404). ResourceRemovedError: If the dataset is no longer available (410). PermissionError: If the dataset's terms have not been agreed to or the dataset is @@ -173,13 +206,16 @@ def download_sample( dataset_details = get_dataset_details(dataset_id) - return _download_sample( + archive_path = _download_sample( dataset_id=dataset_details.id, download_directory=download_directory, show_progress=show_progress, overwrite_existing=overwrite_existing, download_source=DOWNLOAD_SOURCE_SAMPLE, ) + if not extract: + return archive_path + return _extract_beside(archive_path, overwrite_existing, overwrite_extracted) def _download_archive( @@ -189,6 +225,8 @@ def _download_archive( overwrite_existing: bool, enable_logging: bool, download_source: str, + extract: bool = False, + overwrite_extracted: bool = False, ) -> Path: """Shared body of `download_dataset` and its deprecated alias; the *download_source* sent with the request tells them apart.""" @@ -197,7 +235,7 @@ def _download_archive( dataset_details = get_dataset_details(dataset_id) - return _download_dataset( + archive_path = _download_dataset( dataset_id=dataset_details.id, archive_filename=_require_archive_filename(dataset_details), download_directory=download_directory, @@ -205,6 +243,25 @@ def _download_archive( overwrite_existing=overwrite_existing, download_source=download_source, ) + if not extract: + return archive_path + return _extract_beside(archive_path, overwrite_existing, overwrite_extracted) + + +def _extract_beside( + archive_path: Path, overwrite_existing: bool, overwrite_extracted: bool +) -> Path: + """ + Unpack an archive into a directory beside it, named after the archive without its + suffix. The one place that decides when an existing directory is unpacked again: on + request (`overwrite_extracted`), or because the archive itself was downloaded again + (`overwrite_existing`) and the directory may no longer match it. + """ + return _extract_archive( + archive_path=archive_path, + dest_dir=archive_path.parent, + overwrite_extracted=overwrite_existing or overwrite_extracted, + ) # Added these two overload typing declarations in order to accurately type check @@ -255,7 +312,8 @@ def load_dataset( unless `overwrite_existing=True`. If there is a directory with the same name as the archive file without the suffix extension, we assume - it has already been extracted, and it will not be re-extracted unless `overwrite_extracted=True`. + it has already been extracted, and it will not be re-extracted unless `overwrite_extracted=True` + or the archive was downloaded again with `overwrite_existing=True`. Uses the dataset schema to determine the loading strategy. @@ -284,6 +342,8 @@ def load_dataset( is invalid. MissingDependencyError: If `return_format="hf"` and the HuggingFace `datasets` library is not installed. + FileExistsError: If a file, rather than a directory, sits where the archive + would be extracted. FileNotFoundError: If the dataset does not exist (404). ResourceRemovedError: If the dataset is no longer available (410). PermissionError: If access is denied (403) or download directory is not writable. @@ -326,12 +386,7 @@ def load_dataset( download_source=DOWNLOAD_SOURCE_LOAD, ) - base_dir = archive_path.parent - extract_dir = _extract_archive( - archive_path=archive_path, - dest_dir=base_dir, - overwrite_extracted=overwrite_extracted, - ) + extract_dir = _extract_beside(archive_path, overwrite_existing, overwrite_extracted) schema = _resolve_schema(_id, extract_dir, archive_checksum) df = _load_dataset_from_schema(schema, extract_dir) diff --git a/tests/test_archive_utils.py b/tests/test_archive_utils.py index 30e688e..ddd05a1 100644 --- a/tests/test_archive_utils.py +++ b/tests/test_archive_utils.py @@ -3,7 +3,7 @@ import pytest -from datacollective.archive_utils import _extract_archive +from datacollective.archive_utils import _extract_archive, _extracted_dir_name def _make_tar_gz(tmp_path: Path, name: str = "sample.tar.gz") -> Path: @@ -54,3 +54,96 @@ def test_extract_archive_rejects_non_tar_gz(tmp_path: Path) -> None: with pytest.raises(ValueError, match=r"\.tar\.gz or \.tgz"): _extract_archive(archive, tmp_path, overwrite_extracted=False) + + +def test_extract_archive_interrupted_leaves_no_finished_directory( + tmp_path: Path, +) -> None: + archive = _make_tar_gz(tmp_path) + good = archive.read_bytes() + archive.write_bytes(good[: len(good) // 2]) # truncated, as if cut off mid-download + dest = tmp_path / "out" + + with pytest.raises((tarfile.TarError, EOFError)): + _extract_archive(archive, dest, overwrite_extracted=False) + + # Nothing that looks finished, so the next call extracts rather than skipping + assert not (dest / "sample").exists() + + archive.write_bytes(good) + target = _extract_archive(archive, dest, overwrite_extracted=False) + + assert (target / "data.txt").read_text() == "hello" + assert sorted(p.name for p in dest.iterdir()) == ["sample"] + + +def test_extract_archive_overwrite_replaces_directory(tmp_path: Path) -> None: + archive = _make_tar_gz(tmp_path) + stale = tmp_path / "out" / "sample" + stale.mkdir(parents=True) + (stale / "old.txt").write_text("stale") + + target = _extract_archive(archive, tmp_path / "out", overwrite_extracted=True) + + assert sorted(p.name for p in target.iterdir()) == ["data.txt"] + + +def test_extract_archive_failed_overwrite_keeps_existing_directory( + tmp_path: Path, +) -> None: + archive = _make_tar_gz(tmp_path) + good = archive.read_bytes() + dest = tmp_path / "out" + existing = _extract_archive(archive, dest, overwrite_extracted=False) + archive.write_bytes(good[: len(good) // 2]) # the re-downloaded archive is corrupt + + with pytest.raises((tarfile.TarError, EOFError)): + _extract_archive(archive, dest, overwrite_extracted=True) + + # The good directory is untouched and the failed attempt left nothing behind. + assert (existing / "data.txt").read_text() == "hello" + assert sorted(p.name for p in dest.iterdir()) == ["sample"] + + +def test_extract_archive_failure_removes_partial_directory(tmp_path: Path) -> None: + archive = _make_tar_gz(tmp_path) + good = archive.read_bytes() + archive.write_bytes(good[: len(good) // 2]) + dest = tmp_path / "out" + + with pytest.raises((tarfile.TarError, EOFError)): + _extract_archive(archive, dest, overwrite_extracted=False) + + assert not (dest / "sample.partial").exists() + + +@pytest.mark.parametrize("overwrite_extracted", [False, True]) +@pytest.mark.parametrize("stray_name", ["sample", "sample.partial"]) +def test_extract_archive_refuses_to_touch_stray_file( + tmp_path: Path, stray_name: str, overwrite_extracted: bool +) -> None: + archive = _make_tar_gz(tmp_path) + dest = tmp_path / "out" + dest.mkdir() + stray = dest / stray_name + stray.write_text("mine") + + with pytest.raises(FileExistsError, match="is not a directory"): + _extract_archive(archive, dest, overwrite_extracted=overwrite_extracted) + + # The stray file is left exactly as it was and nothing else was created. + assert stray.read_text() == "mine" + assert sorted(p.name for p in dest.iterdir()) == [stray_name] + + +@pytest.mark.parametrize( + ("archive_name", "expected"), + [ + ("dataset.tar.gz", "dataset"), + ("dataset.tgz", "dataset"), + ("dataset.zip", None), + (".tar.gz", None), + ], +) +def test_extracted_dir_name(archive_name: str, expected: str | None) -> None: + assert _extracted_dir_name(archive_name) == expected diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 1461884..8b575df 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -103,3 +103,134 @@ def fake_download_sample(**kwargs: object) -> Path: "download_source": "download_sample", } ] + + +# A re-downloaded archive is always unpacked afresh; otherwise only on request. +_OVERWRITE_CASES = [ + (False, False, False), + (True, False, True), + (False, True, True), + (True, True, True), +] + + +@pytest.mark.parametrize( + ("overwrite_existing", "overwrite_extracted", "expected"), _OVERWRITE_CASES +) +def test_download_dataset_extract_returns_extracted_directory( + monkeypatch: MonkeyPatch, + overwrite_existing: bool, + overwrite_extracted: bool, + expected: bool, +) -> None: + calls = _capture_downloads(monkeypatch) + extracted: list[dict[str, object]] = [] + + def fake_extract_archive(**kwargs: object) -> Path: + extracted.append(kwargs) + return Path("archive") + + monkeypatch.setattr(datasets, "_extract_archive", fake_extract_archive) + + result = datasets.download_dataset( + "ds", + extract=True, + overwrite_existing=overwrite_existing, + overwrite_extracted=overwrite_extracted, + ) + + assert result == Path("archive") + # Re-extracting never triggers a re-download on its own. + assert calls[0]["overwrite_existing"] is overwrite_existing + assert extracted == [ + { + "archive_path": Path("archive.tar.gz"), + "dest_dir": Path("."), + "overwrite_extracted": expected, + } + ] + + +@pytest.mark.parametrize( + ("overwrite_existing", "overwrite_extracted", "expected"), _OVERWRITE_CASES +) +def test_load_dataset_re_extracts_a_re_downloaded_archive( + monkeypatch: MonkeyPatch, + overwrite_existing: bool, + overwrite_extracted: bool, + expected: bool, +) -> None: + _capture_downloads(monkeypatch) + monkeypatch.setattr(datasets, "_get_dataset_schema", lambda _id: object()) + monkeypatch.setattr(datasets, "_resolve_schema", lambda *args: object()) + monkeypatch.setattr( + datasets, "_load_dataset_from_schema", lambda schema, extract_dir: extract_dir + ) + extracted: list[dict[str, object]] = [] + + def fake_extract_archive(**kwargs: object) -> Path: + extracted.append(kwargs) + return Path("archive") + + monkeypatch.setattr(datasets, "_extract_archive", fake_extract_archive) + + result = datasets.load_dataset( + "ds", + overwrite_existing=overwrite_existing, + overwrite_extracted=overwrite_extracted, + ) + + assert result == Path("archive") + assert extracted == [ + { + "archive_path": Path("archive.tar.gz"), + "dest_dir": Path("."), + "overwrite_extracted": expected, + } + ] + + +@pytest.mark.parametrize( + ("overwrite_existing", "overwrite_extracted", "expected"), _OVERWRITE_CASES +) +def test_download_sample_extract_returns_extracted_directory( + monkeypatch: MonkeyPatch, + overwrite_existing: bool, + overwrite_extracted: bool, + expected: bool, +) -> None: + sample = Path("samples") / "ds" / "sample.tar.gz" + calls: list[dict[str, object]] = [] + + def fake_download_sample(**kwargs: object) -> Path: + calls.append(kwargs) + return sample + + monkeypatch.setattr( + datasets, "get_dataset_details", lambda dataset_id: DatasetDetails(id="ds") + ) + monkeypatch.setattr(datasets, "_download_sample", fake_download_sample) + extracted: list[dict[str, object]] = [] + + def fake_extract_archive(**kwargs: object) -> Path: + extracted.append(kwargs) + return sample.parent / "sample" + + monkeypatch.setattr(datasets, "_extract_archive", fake_extract_archive) + + result = datasets.download_sample( + "ds", + extract=True, + overwrite_existing=overwrite_existing, + overwrite_extracted=overwrite_extracted, + ) + + assert result == sample.parent / "sample" + assert calls[0]["overwrite_existing"] is overwrite_existing + assert extracted == [ + { + "archive_path": sample, + "dest_dir": sample.parent, + "overwrite_extracted": expected, + } + ]