From a5dc908c8c268b2fbf2ae32ff34f12fd4452be0b Mon Sep 17 00:00:00 2001 From: Aaron Date: Thu, 15 Oct 2026 20:12:34 +0100 Subject: [PATCH 1/8] Add download_sample for a dataset's public sample The platform serves a dataset's sample from POST /datasets/{id}/sample-download, which the SDK had no call for. download_sample mirrors download_dataset: same arguments, resolves a slug to the dataset id, skips the request when the sample is already on disk, and reports its own download source. Two differences follow from the endpoint. The sample's filename is only known from the response, so samples are kept in samples// under the download directory, and a finished file there means no request is made. The response carries no checksum, so an interrupted sample download starts again rather than resuming. The overwrite, resume, download and rename steps move into _run_download_plan, and response parsing into _build_download_plan, so both downloads share them. --- src/datacollective/__init__.py | 2 + src/datacollective/datasets.py | 57 +++++++++++ src/datacollective/download.py | 171 +++++++++++++++++++++++++++++++++ tests/test_datasets.py | 30 ++++++ tests/test_download.py | 159 ++++++++++++++++++++++++++++++ tests/test_openapi_contract.py | 7 ++ 6 files changed, 426 insertions(+) diff --git a/src/datacollective/__init__.py b/src/datacollective/__init__.py index 065b8a4..46a4d97 100644 --- a/src/datacollective/__init__.py +++ b/src/datacollective/__init__.py @@ -2,6 +2,7 @@ from datacollective.datasets import ( download_dataset, + download_sample, get_dataset_details, list_dataset_filters, list_datasets, @@ -33,6 +34,7 @@ __all__ = [ "download_dataset", + "download_sample", "save_dataset_to_disk", "load_dataset", "get_dataset_details", diff --git a/src/datacollective/datasets.py b/src/datacollective/datasets.py index a2d1597..78a6479 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -23,8 +23,10 @@ from datacollective.download import ( DOWNLOAD_SOURCE_DOWNLOAD, DOWNLOAD_SOURCE_LOAD, + DOWNLOAD_SOURCE_SAMPLE, DOWNLOAD_SOURCE_SAVE, _download_dataset, + _download_sample, ) from datacollective.hf_utils import _convert_to_hf, _require_datasets from datacollective.logging_utils import ( @@ -123,6 +125,61 @@ def download_dataset( ) +def download_sample( + dataset_id: str, + download_directory: str | None = None, + show_progress: bool = True, + overwrite_existing: bool = False, + enable_logging: bool = False, +) -> Path: + """ + Download the dataset's sample archive to a local directory and return the archive path. + The sample is a small excerpt of the dataset that its publisher provides for evaluation. + + The archive is saved under `samples//` inside the download directory, apart + from the full dataset archive. Skips download if the sample is already there (unless + `overwrite_existing=True`, which also picks up a sample the publisher has since replaced). + + 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. + + Requires an API key whose organization has agreed to the dataset's terms on the MDC + platform. Buying the dataset is not required. + + Args: + dataset_id: The dataset ID (as shown in MDC platform) or slug. + download_directory: Directory under which to save the sample archive. + 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. + enable_logging: Whether to enable SDK logging to console and a local log file. + + Returns: + Path to the downloaded sample archive. + + Raises: + ValueError: If dataset_id is empty. + 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 + private to another organization (403), or download directory is not writable. + RuntimeError: If rate limit is exceeded (429) or unexpected response format. + requests.HTTPError: For other non-2xx responses. + """ + _enable_logging(enable_logging) + logger.info(f"Downloading sample of dataset {dataset_id}") + + dataset_details = get_dataset_details(dataset_id) + + return _download_sample( + dataset_id=dataset_details.id, + download_directory=download_directory, + show_progress=show_progress, + overwrite_existing=overwrite_existing, + download_source=DOWNLOAD_SOURCE_SAMPLE, + ) + + def _download_archive( dataset_id: str, download_directory: str | None, diff --git a/src/datacollective/download.py b/src/datacollective/download.py index e8c9b01..4aca182 100644 --- a/src/datacollective/download.py +++ b/src/datacollective/download.py @@ -20,6 +20,11 @@ DOWNLOAD_SOURCE_DOWNLOAD = "download_dataset" DOWNLOAD_SOURCE_SAVE = "save_dataset_to_disk" DOWNLOAD_SOURCE_LOAD = "load_dataset" +DOWNLOAD_SOURCE_SAMPLE = "download_sample" + +# Samples live under /samples//, apart from full archives. +SAMPLES_DIRNAME = "samples" +_PARTIAL_SUFFIXES = (".part", ".checksum") class DownloadPlan(NonEmptyStrModel): @@ -80,6 +85,102 @@ def _download_dataset( download_source=download_source, ) + return _run_download_plan( + download_plan=download_plan, + show_progress=show_progress, + overwrite_existing=overwrite_existing, + download_source=download_source, + ) + + +def _download_sample( + dataset_id: str, + download_directory: str | None, + show_progress: bool, + overwrite_existing: bool, + download_source: str | None = None, +) -> Path: + """ + Download a dataset's sample archive from MDC to a local directory. + + Follows the same flow as `_download_dataset`, with two differences: + + - The sample's filename is only known once the API is called, so samples are kept in a + directory of their own, `/samples//`. A finished file + there means the sample is already downloaded, and the API is not called. + - The API returns no checksum for a sample, so an interrupted download cannot be matched + to the current sample and starts again instead of resuming. + + Args: + dataset_id: The unique dataset ID. + download_directory: The base directory to download into. + show_progress: Show a progress bar when downloading the sample. + overwrite_existing: If True, download the sample again and replace any existing one. + download_source: Optional context appended to the User-Agent for download analytics. + Returns: + The full path to the downloaded sample archive. + """ + sample_dir = ( + _resolve_download_dir(download_directory) / SAMPLES_DIRNAME / dataset_id + ) + sample_dir.mkdir(parents=True, exist_ok=True) + + existing = _find_downloaded_sample(sample_dir) + if existing and not overwrite_existing: + logger.info(f"Skipping download. Sample archive already exists at `{existing}`") + return existing + + download_plan = _get_sample_download_plan( + dataset_id=dataset_id, + sample_dir=sample_dir, + download_source=download_source, + ) + + if download_plan.tmp_filepath.exists(): + logger.info("Samples cannot be resumed. Starting the sample download again...") + _cleanup_partial_download(download_plan) + + target_filepath = _run_download_plan( + download_plan=download_plan, + show_progress=show_progress, + overwrite_existing=overwrite_existing, + download_source=download_source, + ) + + # Keep one sample per dataset: remove any earlier one the publisher has since replaced, + # along with its leftover partial files. + for other in sample_dir.iterdir(): + if ( + other != target_filepath + and other.is_file() + and not other.name.startswith(".") + ): + other.unlink() + + return target_filepath + + +def _run_download_plan( + download_plan: DownloadPlan, + show_progress: bool, + overwrite_existing: bool, + download_source: str | None = None, +) -> Path: + """ + Carry out a download plan: clean up when overwriting, resume where possible, download to + the temporary file, then move it into place. + + Args: + download_plan: The DownloadPlan object with download details. + show_progress: Whether to show a progress bar during download. + overwrite_existing: If True, remove any existing complete or partial download first. + download_source: Optional context appended to the User-Agent for download analytics. + + Returns: + The full path to the downloaded archive. + """ + target_filepath = download_plan.target_filepath + # If overwriting, clean up any existing complete or partial download files if overwrite_existing: logger.info( @@ -138,6 +239,63 @@ def _get_download_plan( ) payload: dict[str, Any] = resp.json() + return _build_download_plan(payload, target_filepath) + + +def _get_sample_download_plan( + dataset_id: str, + sample_dir: Path, + download_source: str | None = None, +) -> DownloadPlan: + """ + Send a POST request to the API to receive the download details for a dataset's sample. + + Args: + dataset_id: The dataset ID (as shown in MDC platform). + sample_dir: Resolved, absolute directory to save the sample in. + download_source: Optional context appended to the User-Agent for download analytics. + + Returns: + a DownloadPlan object + + Raises: + FileNotFoundError: If the dataset does not exist or publishes no sample (404). + ResourceRemovedError: If the dataset is no longer available (410). + PermissionError: If access is denied (403). + RuntimeError: If rate limit is exceeded (429) or unexpected response format. + requests.HTTPError: For other non-2xx responses. + """ + sample_url = f"{_get_api_url()}/datasets/{dataset_id}/sample-download" + resp = _send_api_request( + method="POST", url=sample_url, source_function=download_source + ) + + payload: dict[str, Any] = resp.json() + filename = str(payload.get("filename") or "") + # The filename becomes a path on disk, so accept a bare name only. Hidden names are + # refused too, since they are skipped when looking for an existing sample. + if Path(filename).name != filename or not filename or filename.startswith("."): + raise RuntimeError(f"Unexpected sample filename in response: {filename!r}") + + return _build_download_plan(payload, sample_dir / filename) + + +def _build_download_plan( + payload: dict[str, Any], target_filepath: Path +) -> DownloadPlan: + """ + Build a DownloadPlan from a download response of the MDC API. + + Args: + payload: The JSON body of the response. + target_filepath: Resolved, absolute full file path to save the download. + + Returns: + a DownloadPlan object + + Raises: + RuntimeError: If the response lacks a download URL or size. + """ download_url = payload.get("downloadUrl") size_bytes = payload.get("sizeBytes") checksum = payload.get("checksum") @@ -288,6 +446,19 @@ def _resolve_download_dir(download_directory: str | None) -> Path: return p +def _find_downloaded_sample(sample_dir: Path) -> Path | None: + """Return the finished sample archive in *sample_dir*, or None if there is none. + Hidden files, such as the `.DS_Store` macOS leaves behind, are not samples.""" + finished = [ + p + for p in sample_dir.iterdir() + if p.is_file() + and not p.name.startswith(".") + and not p.name.endswith(_PARTIAL_SUFFIXES) + ] + return max(finished, key=lambda p: p.stat().st_mtime, default=None) + + def _get_checksum_filepath(target_filepath: Path) -> Path: """Return the path to the .checksum file for a given target file.""" return target_filepath.with_suffix(target_filepath.suffix + ".checksum") diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 7b71ac7..1461884 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -73,3 +73,33 @@ def test_save_dataset_to_disk_warns_and_reports_its_own_source( "download_source": "save_dataset_to_disk", } ] + + +def test_download_sample_resolves_id_and_reports_its_own_source( + monkeypatch: MonkeyPatch, +) -> None: + calls: list[dict[str, object]] = [] + + def fake_download_sample(**kwargs: object) -> Path: + calls.append(kwargs) + return Path("sample.tar.gz") + + monkeypatch.setattr( + datasets, + "get_dataset_details", + lambda dataset_id: DatasetDetails(id="canonical-id"), + ) + monkeypatch.setattr(datasets, "_download_sample", fake_download_sample) + + result = datasets.download_sample("a-dataset-slug", download_directory="dir") + + assert result == Path("sample.tar.gz") + assert calls == [ + { + "dataset_id": "canonical-id", + "download_directory": "dir", + "show_progress": True, + "overwrite_existing": False, + "download_source": "download_sample", + } + ] diff --git a/tests/test_download.py b/tests/test_download.py index f19fc0c..3fb4861 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -3,8 +3,11 @@ import logging from typing import Any +import pytest + from datacollective.download import ( DownloadPlan, + _download_sample, _get_checksum_filepath, _determine_resume_state, _resolve_download_dir, @@ -229,3 +232,159 @@ def fake_send_api_request(**kwargs: Any) -> FakeResponse: assert captured["url"] == plan.download_url assert captured["source_function"] == "download_dataset" assert captured["include_auth_headers"] is False + + +""" +Sample downloads +""" + + +class _FakeSampleApi: + """Stands in for `_send_api_request`: answers the sample-download POST with *payload* + and the presigned GET with *content*, recording every call.""" + + def __init__(self, payload: dict[str, Any], content: bytes = b"sample") -> None: + self.payload = payload + self.content = content + self.calls: list[dict[str, Any]] = [] + + def __call__(self, **kwargs: Any) -> Any: + self.calls.append(kwargs) + api = self + + class FakeResponse: + def json(self) -> dict[str, Any]: + return api.payload + + def __enter__(self) -> "FakeResponse": + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + return None + + def iter_content(self, chunk_size: int): + yield api.content + + return FakeResponse() + + +def _sample_payload(filename: str = "sample.tar.gz") -> dict[str, Any]: + return { + "downloadUrl": "https://example.com/sample", + "filename": filename, + "sizeBytes": "6", + "contentType": "application/gzip", + "expiresAt": "2026-08-27T12:00:00.000Z", + } + + +def test_download_sample_saves_under_samples_dir(tmp_path: Path, monkeypatch) -> None: + api = _FakeSampleApi(_sample_payload()) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download_sample( + "ds", + str(tmp_path), + show_progress=False, + overwrite_existing=False, + download_source="download_sample", + ) + + assert result == tmp_path / "samples" / "ds" / "sample.tar.gz" + assert result.read_bytes() == b"sample" + assert sorted(p.name for p in result.parent.iterdir()) == ["sample.tar.gz"] + + post, get = api.calls + assert post["method"] == "POST" + assert post["url"].endswith("/datasets/ds/sample-download") + assert post["source_function"] == "download_sample" + assert post.get("include_auth_headers", True) is True + assert get["method"] == "GET" + assert get["url"] == "https://example.com/sample" + assert get["include_auth_headers"] is False + + +def test_download_sample_skips_api_when_already_downloaded( + tmp_path: Path, monkeypatch +) -> None: + existing = tmp_path / "samples" / "ds" / "sample.tar.gz" + existing.parent.mkdir(parents=True) + existing.write_bytes(b"cached") + api = _FakeSampleApi(_sample_payload()) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert result == existing + assert result.read_bytes() == b"cached" + assert api.calls == [] + + +def test_download_sample_ignores_hidden_files(tmp_path: Path, monkeypatch) -> None: + sample_dir = tmp_path / "samples" / "ds" + sample_dir.mkdir(parents=True) + (sample_dir / ".DS_Store").write_bytes(b"finder") + api = _FakeSampleApi(_sample_payload()) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert result.name == "sample.tar.gz" + assert (sample_dir / ".DS_Store").exists() + + +def test_download_sample_restarts_interrupted_download( + tmp_path: Path, monkeypatch, caplog +) -> None: + partial = tmp_path / "samples" / "ds" / "sample.tar.gz.part" + partial.parent.mkdir(parents=True) + partial.write_bytes(b"stale") + api = _FakeSampleApi(_sample_payload()) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + with caplog.at_level(logging.INFO, logger="datacollective.download"): + result = _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert result.read_bytes() == b"sample" + assert not partial.exists() + assert "Samples cannot be resumed" in caplog.text + + +def test_download_sample_overwrite_replaces_previous_sample( + tmp_path: Path, monkeypatch +) -> None: + old = tmp_path / "samples" / "ds" / "old-sample.tar.gz" + old.parent.mkdir(parents=True) + old.write_bytes(b"old") + api = _FakeSampleApi(_sample_payload("new-sample.tar.gz")) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + result = _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=True + ) + + assert result.name == "new-sample.tar.gz" + assert sorted(p.name for p in result.parent.iterdir()) == ["new-sample.tar.gz"] + + +@pytest.mark.parametrize( + "filename", ["../escape.tar.gz", "nested/sample.tar.gz", "", ".hidden.tar.gz"] +) +def test_download_sample_rejects_unsafe_filename( + tmp_path: Path, monkeypatch, filename: str +) -> None: + api = _FakeSampleApi(_sample_payload(filename)) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + with pytest.raises(RuntimeError, match="Unexpected sample filename"): + _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert len(api.calls) == 1 # nothing downloaded diff --git a/tests/test_openapi_contract.py b/tests/test_openapi_contract.py index da70906..31b85ba 100644 --- a/tests/test_openapi_contract.py +++ b/tests/test_openapi_contract.py @@ -50,6 +50,7 @@ SDK_OPERATIONS = { ("get", "/datasets/{datasetId}"), ("post", "/datasets/{datasetId}/download"), + ("post", "/datasets/{datasetId}/sample-download"), ("post", "/submissions"), ("get", "/submissions/{submissionId}"), ("patch", "/submissions/{submissionId}"), @@ -321,6 +322,12 @@ def test_download_session_fields(schemas: dict[str, Any]) -> None: assert "checksum" in _properties(schema, schemas) +def test_sample_download_fields(schemas: dict[str, Any]) -> None: + schema = _resolve(schemas["SampleDownload"], schemas) + # download._get_sample_download_plan reads these keys and requires all three. + assert {"downloadUrl", "sizeBytes", "filename"} <= _required(schema) + + """ Check upload session spec """ From 984cb68152248032dfded02db0d151c711275760 Mon Sep 17 00:00:00 2001 From: Aaron Date: Thu, 15 Oct 2026 20:30:24 +0100 Subject: [PATCH 2/8] Create a sample's directory only once there is a sample to save A refused sample request, such as one for a dataset whose terms have not been agreed to, left an empty samples// directory behind. --- src/datacollective/download.py | 6 +++++- tests/test_download.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/datacollective/download.py b/src/datacollective/download.py index 4aca182..0f04ff8 100644 --- a/src/datacollective/download.py +++ b/src/datacollective/download.py @@ -123,7 +123,6 @@ def _download_sample( sample_dir = ( _resolve_download_dir(download_directory) / SAMPLES_DIRNAME / dataset_id ) - sample_dir.mkdir(parents=True, exist_ok=True) existing = _find_downloaded_sample(sample_dir) if existing and not overwrite_existing: @@ -136,6 +135,9 @@ def _download_sample( download_source=download_source, ) + # Created only now, so a refused request leaves no empty directory behind. + sample_dir.mkdir(parents=True, exist_ok=True) + if download_plan.tmp_filepath.exists(): logger.info("Samples cannot be resumed. Starting the sample download again...") _cleanup_partial_download(download_plan) @@ -449,6 +451,8 @@ def _resolve_download_dir(download_directory: str | None) -> Path: def _find_downloaded_sample(sample_dir: Path) -> Path | None: """Return the finished sample archive in *sample_dir*, or None if there is none. Hidden files, such as the `.DS_Store` macOS leaves behind, are not samples.""" + if not sample_dir.is_dir(): + return None finished = [ p for p in sample_dir.iterdir() diff --git a/tests/test_download.py b/tests/test_download.py index 3fb4861..4e51c16 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -388,3 +388,19 @@ def test_download_sample_rejects_unsafe_filename( ) assert len(api.calls) == 1 # nothing downloaded + + +def test_download_sample_refused_leaves_no_directory( + tmp_path: Path, monkeypatch +) -> None: + def refuse(**kwargs: Any) -> Any: + raise PermissionError("You must agree to the terms of use") + + monkeypatch.setattr("datacollective.download._send_api_request", refuse) + + with pytest.raises(PermissionError): + _download_sample( + "ds", str(tmp_path), show_progress=False, overwrite_existing=False + ) + + assert not (tmp_path / "samples").exists() From f516c87b3154f5688f5a0fd5ea6e48a4073dac99 Mon Sep 17 00:00:00 2001 From: Aaron Date: Thu, 15 Oct 2026 20:35:08 +0100 Subject: [PATCH 3/8] Add extract to download_dataset and download_sample Extraction was only reachable through load_dataset, which needs the dataset in the schema registry. With extract=True, download_dataset and download_sample also unpack the archive into a directory beside it and return that directory. It is the directory load_dataset extracts to, so the two never unpack the same archive twice. An existing directory is reused unless overwrite_existing=True, which re-downloads and so re-extracts. _extract_archive now extracts under a temporary name and renames once complete. Before, it created the target directory first, so an interrupted extraction left a directory that the next call took as finished and skipped. load_dataset gets the fix too. Replacing a sample also removes the directory the old one was extracted into. --- src/datacollective/archive_utils.py | 28 +++++++++--- src/datacollective/datasets.py | 53 ++++++++++++++++++---- src/datacollective/download.py | 7 ++- tests/test_archive_utils.py | 47 +++++++++++++++++++- tests/test_datasets.py | 68 +++++++++++++++++++++++++++++ tests/test_download.py | 15 +++++++ 6 files changed, 203 insertions(+), 15 deletions(-) diff --git a/src/datacollective/archive_utils.py b/src/datacollective/archive_utils.py index 5521baa..74d377a 100644 --- a/src/datacollective/archive_utils.py +++ b/src/datacollective/archive_utils.py @@ -17,6 +17,9 @@ 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. + Args: archive_path: Path to the archive file. dest_dir: Directory where to extract the contents. @@ -27,15 +30,15 @@ def _extract_archive( Raises: ValueError: If the archive is not a `.tar.gz` / `.tgz` file. """ - 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) + target = dest_dir / dir_name if target.exists(): if not overwrite_extracted: logger.info( @@ -46,7 +49,22 @@ def _extract_archive( logger.info(f"Overwriting existing extracted directory: `{target}`") shutil.rmtree(target) - target.mkdir(parents=True, exist_ok=True) + # Left over from an interrupted extraction; start again. + partial = target.with_name(target.name + ".partial") + if partial.exists(): + shutil.rmtree(partial) + + partial.mkdir(parents=True) with tarfile.open(archive_path, "r:gz") as tf: - tf.extractall(path=target, filter="data") + tf.extractall(path=partial, filter="data") + partial.rename(target) 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 78a6479..72edb0f 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -85,6 +85,7 @@ def download_dataset( show_progress: bool = True, overwrite_existing: bool = False, enable_logging: bool = False, + extract: bool = False, ) -> Path: """ Download the dataset archive to a local directory and return the archive path. @@ -93,6 +94,11 @@ 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_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 +107,18 @@ 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, and with + `extract=True` the extracted directory 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. 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. 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 +132,7 @@ def download_dataset( overwrite_existing, enable_logging, download_source=DOWNLOAD_SOURCE_DOWNLOAD, + extract=extract, ) @@ -131,6 +142,7 @@ def download_sample( show_progress: bool = True, overwrite_existing: bool = False, enable_logging: bool = False, + extract: bool = False, ) -> Path: """ Download the dataset's sample archive to a local directory and return the archive path. @@ -143,6 +155,10 @@ 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_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. @@ -151,14 +167,18 @@ def download_sample( download_directory: Directory under which to save the sample archive. 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. + overwrite_existing: Whether to download the sample again, replacing the existing one + and, with `extract=True`, its extracted directory. 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. 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 archive is not a + `.tar.gz` / `.tgz` file. 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 @@ -171,13 +191,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) def _download_archive( @@ -187,6 +210,7 @@ def _download_archive( overwrite_existing: bool, enable_logging: bool, download_source: str, + extract: bool = False, ) -> Path: """Shared body of `download_dataset` and its deprecated alias; the *download_source* sent with the request tells them apart.""" @@ -195,7 +219,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, @@ -203,6 +227,19 @@ def _download_archive( overwrite_existing=overwrite_existing, download_source=download_source, ) + if not extract: + return archive_path + return _extract_beside(archive_path, overwrite_existing) + + +def _extract_beside(archive_path: Path, overwrite_existing: bool) -> Path: + """Unpack an archive into a directory beside it, as `load_dataset` does. A + re-downloaded archive (`overwrite_existing`) is unpacked afresh.""" + return _extract_archive( + archive_path=archive_path, + dest_dir=archive_path.parent, + overwrite_extracted=overwrite_existing, + ) # Added these two overload typing declarations in order to accurately type check diff --git a/src/datacollective/download.py b/src/datacollective/download.py index 0f04ff8..95a0f72 100644 --- a/src/datacollective/download.py +++ b/src/datacollective/download.py @@ -1,10 +1,12 @@ import os +import shutil from pathlib import Path from typing import Any from fox_progress_bar import ProgressBar from pydantic import Field +from datacollective.archive_utils import _extracted_dir_name from datacollective.api_utils import ( ENV_DOWNLOAD_PATH, HTTP_TIMEOUT, @@ -150,7 +152,7 @@ def _download_sample( ) # Keep one sample per dataset: remove any earlier one the publisher has since replaced, - # along with its leftover partial files. + # along with its leftover partial files and the directory it was extracted into. for other in sample_dir.iterdir(): if ( other != target_filepath @@ -158,6 +160,9 @@ def _download_sample( and not other.name.startswith(".") ): other.unlink() + stale_dir_name = _extracted_dir_name(other.name) + if stale_dir_name and (sample_dir / stale_dir_name).is_dir(): + shutil.rmtree(sample_dir / stale_dir_name) return target_filepath diff --git a/tests/test_archive_utils.py b/tests/test_archive_utils.py index 30e688e..e1a808a 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,48 @@ 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"] + + +@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..ba6e616 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -103,3 +103,71 @@ def fake_download_sample(**kwargs: object) -> Path: "download_source": "download_sample", } ] + + +@pytest.mark.parametrize("overwrite_existing", [False, True]) +def test_download_dataset_extract_returns_extracted_directory( + monkeypatch: MonkeyPatch, overwrite_existing: bool +) -> None: + _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 + ) + + assert result == Path("archive") + assert extracted == [ + { + "archive_path": Path("archive.tar.gz"), + "dest_dir": Path("."), + "overwrite_extracted": overwrite_existing, + } + ] + + +def test_download_dataset_without_extract_does_not_extract( + monkeypatch: MonkeyPatch, +) -> None: + _capture_downloads(monkeypatch) + monkeypatch.setattr( + datasets, + "_extract_archive", + lambda **kwargs: pytest.fail("extracted without extract=True"), + ) + + assert datasets.download_dataset("ds") == Path("archive.tar.gz") + + +def test_download_sample_extract_returns_extracted_directory( + monkeypatch: MonkeyPatch, +) -> None: + sample = Path("samples") / "ds" / "sample.tar.gz" + monkeypatch.setattr( + datasets, "get_dataset_details", lambda dataset_id: DatasetDetails(id="ds") + ) + monkeypatch.setattr(datasets, "_download_sample", lambda **kwargs: 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=True) + + assert result == sample.parent / "sample" + assert extracted == [ + { + "archive_path": sample, + "dest_dir": sample.parent, + "overwrite_extracted": True, + } + ] diff --git a/tests/test_download.py b/tests/test_download.py index 4e51c16..839228d 100644 --- a/tests/test_download.py +++ b/tests/test_download.py @@ -373,6 +373,21 @@ def test_download_sample_overwrite_replaces_previous_sample( assert sorted(p.name for p in result.parent.iterdir()) == ["new-sample.tar.gz"] +def test_download_sample_replacing_removes_old_extracted_directory( + tmp_path: Path, monkeypatch +) -> None: + sample_dir = tmp_path / "samples" / "ds" + (sample_dir / "old-sample").mkdir(parents=True) + (sample_dir / "old-sample" / "clip.wav").write_bytes(b"old") + (sample_dir / "old-sample.tar.gz").write_bytes(b"old") + api = _FakeSampleApi(_sample_payload("new-sample.tar.gz")) + monkeypatch.setattr("datacollective.download._send_api_request", api) + + _download_sample("ds", str(tmp_path), show_progress=False, overwrite_existing=True) + + assert sorted(p.name for p in sample_dir.iterdir()) == ["new-sample.tar.gz"] + + @pytest.mark.parametrize( "filename", ["../escape.tar.gz", "nested/sample.tar.gz", "", ".hidden.tar.gz"] ) From d5abbacdc142a47c103f60e7c577cef6dc793f33 Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 19 Oct 2026 16:13:59 +0300 Subject: [PATCH 4/8] Update docs --- docs/index.md | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) 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 From ecdcbfd519799a459f8758b56d325e7dac872807 Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 19 Oct 2026 16:14:05 +0300 Subject: [PATCH 5/8] Update tests --- tests/test_datasets.py | 53 +++++++++++++++++++++++++++++++++++------- 1 file changed, 45 insertions(+), 8 deletions(-) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 26a6dc9..0f2f7b8 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -105,11 +105,25 @@ def fake_download_sample(**kwargs: object) -> Path: ] -@pytest.mark.parametrize("overwrite_existing", [False, True]) +# 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 + monkeypatch: MonkeyPatch, + overwrite_existing: bool, + overwrite_extracted: bool, + expected: bool, ) -> None: - _capture_downloads(monkeypatch) + calls = _capture_downloads(monkeypatch) extracted: list[dict[str, object]] = [] def fake_extract_archive(**kwargs: object) -> Path: @@ -119,27 +133,44 @@ def fake_extract_archive(**kwargs: object) -> Path: monkeypatch.setattr(datasets, "_extract_archive", fake_extract_archive) result = datasets.download_dataset( - "ds", extract=True, overwrite_existing=overwrite_existing + "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": overwrite_existing, + "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", lambda **kwargs: sample) + monkeypatch.setattr(datasets, "_download_sample", fake_download_sample) extracted: list[dict[str, object]] = [] def fake_extract_archive(**kwargs: object) -> Path: @@ -148,13 +179,19 @@ def fake_extract_archive(**kwargs: object) -> Path: monkeypatch.setattr(datasets, "_extract_archive", fake_extract_archive) - result = datasets.download_sample("ds", extract=True, overwrite_existing=True) + 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": True, + "overwrite_extracted": expected, } ] From 210dd1c87e86d565b592cf9c1041aeb5af8f0349 Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 19 Oct 2026 16:14:16 +0300 Subject: [PATCH 6/8] Separate overwrite_extracted arg --- src/datacollective/datasets.py | 39 ++++++++++++++++++++++------------ 1 file changed, 26 insertions(+), 13 deletions(-) diff --git a/src/datacollective/datasets.py b/src/datacollective/datasets.py index b84c98f..ba74027 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -86,6 +86,7 @@ def download_dataset( 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. @@ -96,8 +97,9 @@ def download_dataset( 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_existing=True`. This is the directory - `load_dataset` extracts to, so the two never unpack the same archive twice. + 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. @@ -107,10 +109,12 @@ 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, and with - `extract=True` the extracted directory too. + 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, or to the extracted directory if @@ -133,6 +137,7 @@ def download_dataset( enable_logging, download_source=DOWNLOAD_SOURCE_DOWNLOAD, extract=extract, + overwrite_extracted=overwrite_extracted, ) @@ -143,6 +148,7 @@ def download_sample( 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. @@ -159,7 +165,8 @@ def download_sample( 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_existing=True`. + 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. @@ -170,14 +177,19 @@ 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. 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 @@ -199,7 +211,7 @@ def download_sample( ) if not extract: return archive_path - return _extract_beside(archive_path, overwrite_existing) + return _extract_beside(archive_path, overwrite_existing or overwrite_extracted) def _download_archive( @@ -210,6 +222,7 @@ def _download_archive( 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.""" @@ -228,16 +241,16 @@ def _download_archive( ) if not extract: return archive_path - return _extract_beside(archive_path, overwrite_existing) + return _extract_beside(archive_path, overwrite_existing or overwrite_extracted) -def _extract_beside(archive_path: Path, overwrite_existing: bool) -> Path: - """Unpack an archive into a directory beside it, as `load_dataset` does. A - re-downloaded archive (`overwrite_existing`) is unpacked afresh.""" +def _extract_beside(archive_path: Path, overwrite_extracted: bool) -> Path: + """Unpack an archive into a directory beside it, as `load_dataset` does. Callers pass + `overwrite_extracted=True` when a re-downloaded archive must be unpacked afresh too.""" return _extract_archive( archive_path=archive_path, dest_dir=archive_path.parent, - overwrite_extracted=overwrite_existing, + overwrite_extracted=overwrite_extracted, ) From da105a60f440f246d94ca9c1ad586ac57d549c03 Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 19 Oct 2026 16:19:08 +0300 Subject: [PATCH 7/8] Dont touch stray files --- src/datacollective/archive_utils.py | 13 ++++++++++++- src/datacollective/datasets.py | 6 ++++++ tests/test_archive_utils.py | 19 +++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/datacollective/archive_utils.py b/src/datacollective/archive_utils.py index 74d377a..bcf2747 100644 --- a/src/datacollective/archive_utils.py +++ b/src/datacollective/archive_utils.py @@ -29,6 +29,8 @@ 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`. """ dir_name = _extracted_dir_name(archive_path.name) if dir_name is None: @@ -39,6 +41,16 @@ def _extract_archive( # Extract into a dedicated directory under `dest_dir` named after the archive 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." + ) + if target.exists(): if not overwrite_extracted: logger.info( @@ -50,7 +62,6 @@ def _extract_archive( shutil.rmtree(target) # Left over from an interrupted extraction; start again. - partial = target.with_name(target.name + ".partial") if partial.exists(): shutil.rmtree(partial) diff --git a/src/datacollective/datasets.py b/src/datacollective/datasets.py index ba74027..93b6a2b 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -123,6 +123,8 @@ def download_dataset( Raises: 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. @@ -190,6 +192,8 @@ def download_sample( Raises: 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 @@ -331,6 +335,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. diff --git a/tests/test_archive_utils.py b/tests/test_archive_utils.py index e1a808a..d4a84df 100644 --- a/tests/test_archive_utils.py +++ b/tests/test_archive_utils.py @@ -88,6 +88,25 @@ def test_extract_archive_overwrite_replaces_directory(tmp_path: Path) -> None: assert sorted(p.name for p in target.iterdir()) == ["data.txt"] +@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"), [ From be539c6716939b0ba742aa5126f245e75deba767 Mon Sep 17 00:00:00 2001 From: Kostis-S-Z Date: Mon, 19 Oct 2026 16:29:44 +0300 Subject: [PATCH 8/8] Fix overwriting logic keeping stale extractions --- src/datacollective/archive_utils.py | 33 ++++++++++++++---------- src/datacollective/datasets.py | 28 +++++++++++---------- tests/test_archive_utils.py | 29 +++++++++++++++++++++ tests/test_datasets.py | 39 +++++++++++++++++++++++++++++ 4 files changed, 103 insertions(+), 26 deletions(-) diff --git a/src/datacollective/archive_utils.py b/src/datacollective/archive_utils.py index bcf2747..51722d1 100644 --- a/src/datacollective/archive_utils.py +++ b/src/datacollective/archive_utils.py @@ -18,7 +18,10 @@ def _extract_archive( 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. + 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. @@ -51,24 +54,28 @@ def _extract_archive( "directory. Move or remove it first." ) - if target.exists(): - if not overwrite_extracted: - logger.info( - f"Extracted directory already exists. Skipping extraction: `{target}`" - ) - 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) partial.mkdir(parents=True) - with tarfile.open(archive_path, "r:gz") as tf: - tf.extractall(path=partial, filter="data") - partial.rename(target) + 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 diff --git a/src/datacollective/datasets.py b/src/datacollective/datasets.py index 93b6a2b..eeedc40 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -215,7 +215,7 @@ def download_sample( ) if not extract: return archive_path - return _extract_beside(archive_path, overwrite_existing or overwrite_extracted) + return _extract_beside(archive_path, overwrite_existing, overwrite_extracted) def _download_archive( @@ -245,16 +245,22 @@ def _download_archive( ) if not extract: return archive_path - return _extract_beside(archive_path, overwrite_existing or overwrite_extracted) + return _extract_beside(archive_path, overwrite_existing, overwrite_extracted) -def _extract_beside(archive_path: Path, overwrite_extracted: bool) -> Path: - """Unpack an archive into a directory beside it, as `load_dataset` does. Callers pass - `overwrite_extracted=True` when a re-downloaded archive must be unpacked afresh too.""" +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_extracted, + overwrite_extracted=overwrite_existing or overwrite_extracted, ) @@ -306,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. @@ -379,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 d4a84df..ddd05a1 100644 --- a/tests/test_archive_utils.py +++ b/tests/test_archive_utils.py @@ -88,6 +88,35 @@ def test_extract_archive_overwrite_replaces_directory(tmp_path: Path) -> None: 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( diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 0f2f7b8..8b575df 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -151,6 +151,45 @@ def fake_extract_archive(**kwargs: object) -> Path: ] +@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 )