Add extract to download_dataset and download_sample

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/<dataset-id>/` 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

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)

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,14 +235,33 @@ 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,
         show_progress=show_progress,
         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)