Adding a registry to have the hashes of datasets (restructured for aws s3)

.github/workflows/test.yaml

@@ -42,14 +42,15 @@ jobs:
         id: data-cache
         uses: actions/cache@v4
         with:
-          path: |
-            ~/.cache/squidpy/*.h5ad
-            ~/.cache/squidpy/*.zarr
+          path: ~/.cache/squidpy
           key: data-${{ hashFiles('**/download_data.py') }}
+          restore-keys: |
+            data-
           enableCrossOsArchive: true
 
       - name: Download datasets
-        if: steps.data-cache.outputs.cache-hit != 'true'
+        # Always run to ensure any missing files are downloaded
+        # (restore-keys may provide partial cache)
         run: uvx hatch run data:download
 
   # Get the test environment from hatch as defined in pyproject.toml.
@@ -122,10 +123,10 @@ jobs:
         id: data-cache
         uses: actions/cache@v4
         with:
-          path: |
-            ~/.cache/squidpy/*.h5ad
-            ~/.cache/squidpy/*.zarr
+          path: ~/.cache/squidpy
           key: data-${{ hashFiles('**/download_data.py') }}
+          restore-keys: |
+            data-
           enableCrossOsArchive: true
 
       - name: System dependencies (Linux)
@@ -181,10 +182,10 @@ jobs:
         id: coverage-data-cache
         uses: actions/cache@v4
         with:
-          path: |
-            ~/.cache/squidpy/*.h5ad
-            ~/.cache/squidpy/*.zarr
+          path: ~/.cache/squidpy
           key: data-${{ hashFiles('**/download_data.py') }}
+          restore-keys: |
+            data-
           enableCrossOsArchive: true
 
       - name: System dependencies (Linux)

.scripts/ci/download_data.py

@@ -1,77 +1,81 @@
 #!/usr/bin/env python3
+"""Download datasets to populate CI cache.
+
+This script downloads all datasets that tests might need.
+The downloader handles caching to DEFAULT_CACHE_DIR (~/.cache/squidpy).
+"""
+
 from __future__ import annotations
 
 import argparse
-from pathlib import Path
-from typing import Any
-
-from squidpy.datasets import visium_hne_sdata
 
 _CNT = 0  # increment this when you want to rebuild the CI cache
-_ROOT = Path.home() / ".cache" / "squidpy"
-
-
-def _print_message(func_name: str, path: Path, *, dry_run: bool = False) -> None:
-    prefix = "[DRY RUN]" if dry_run else ""
-    if path.is_file():
-        print(f"{prefix}[Loading]     {func_name:>25} <- {str(path):>25}")
-    else:
-        print(f"{prefix}[Downloading] {func_name:>25} -> {str(path):>25}")
-
-
-def _maybe_download_data(func_name: str, path: Path) -> Any:
-    import squidpy as sq
-
-    try:
-        return getattr(sq.datasets, func_name)(path=path)
-    except Exception as e:  # noqa: BLE001
-        print(f"File {str(path):>25} seems to be corrupted: {e}. Removing and retrying")
-        path.unlink()
-
-        return getattr(sq.datasets, func_name)(path=path)
 
 
 def main(args: argparse.Namespace) -> None:
     from anndata import AnnData
 
     import squidpy as sq
+    from squidpy.datasets._downloader import DEFAULT_CACHE_DIR
+    from squidpy.datasets._registry import get_registry
 
-    all_datasets = sq.datasets._dataset.__all__ + sq.datasets._image.__all__
-    all_extensions = ["h5ad"] * len(sq.datasets._dataset.__all__) + ["tiff"] * len(sq.datasets._image.__all__)
+    registry = get_registry()
+    print(f"Cache directory: {DEFAULT_CACHE_DIR}")
+
+    # Visium samples tested in CI
+    # Add any sample here that's used in tests to ensure it's cached
+    visium_samples_to_cache = [
+        "V1_Mouse_Kidney",
+        "Targeted_Visium_Human_SpinalCord_Neuroscience",
+        "Visium_FFPE_Human_Breast_Cancer",
+    ]
 
     if args.dry_run:
-        for func_name, ext in zip(all_datasets, all_extensions):
-            if func_name == "visium_hne_sdata":
-                ext = "zarr"
-            path = _ROOT / f"{func_name}.{ext}"
-            _print_message(func_name, path, dry_run=True)
+        print("\nWould download:")
+        print(f"  - {len(registry.anndata_datasets)} AnnData datasets")
+        print(f"  - {len(registry.image_datasets)} Image datasets")
+        print(f"  - {len(registry.spatialdata_datasets)} SpatialData datasets")
+        print(f"  - {len(visium_samples_to_cache)} Visium samples")
         return
 
-    # could be parallelized, but on CI it largely does not matter (usually limited to 2 cores + bandwidth limit)
-    for func_name, ext in zip(all_datasets, all_extensions):
-        if func_name == "visium_hne_sdata":
-            ext = "zarr"
-            path = _ROOT / f"{func_name}.{ext}"
-
-            _print_message(func_name, path)
-            obj = visium_hne_sdata(_ROOT)
-
-            assert path.is_dir(), f"Expected a .zarr folder at {path}"
-            continue
-
-        path = _ROOT / f"{func_name}.{ext}"
-        _print_message(func_name, path)
-        obj = _maybe_download_data(func_name, path)
-
-        # we could do without the AnnData check as well (1 less req. in tox.ini), but it's better to be safe
-        assert isinstance(obj, AnnData | sq.im.ImageContainer), type(obj)
-        assert path.is_file(), path
+    # Download AnnData datasets - just call the function, it handles caching
+    print("\nDownloading AnnData datasets...")
+    for name in registry.anndata_datasets:
+        print(f"  {name}")
+        obj = getattr(sq.datasets, name)()
+        assert isinstance(obj, AnnData), f"Expected AnnData, got {type(obj)}"
+
+    # Download image datasets
+    print("\nDownloading image datasets...")
+    for name in registry.image_datasets:
+        print(f"  {name}")
+        obj = getattr(sq.datasets, name)()
+        assert isinstance(obj, sq.im.ImageContainer)
+
+    # Download SpatialData datasets
+    print("\nDownloading SpatialData datasets...")
+    for name in registry.spatialdata_datasets:
+        print(f"  {name}")
+        obj = getattr(sq.datasets, name)()
+        # Returns SpatialData object
+
+    # Download Visium samples (needed for tests)
+    # Include high-res images since tests use include_hires_tiff=True
+    print("\nDownloading Visium samples (with high-res images)...")
+    for sample in visium_samples_to_cache:
+        print(f"  {sample}")
+        obj = sq.datasets.visium(sample, include_hires_tiff=True)
+        assert isinstance(obj, AnnData), f"Expected AnnData, got {type(obj)}"
+
+    print("\nDone!")
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Download data used for tutorials/examples.")
+    parser = argparse.ArgumentParser(description="Download datasets to populate CI cache.")
     parser.add_argument(
-        "--dry-run", action="store_true", help="Do not download any data, just print what would be downloaded."
+        "--dry-run",
+        action="store_true",
+        help="Do not download, just print what would be downloaded.",
     )
 
     main(parser.parse_args())

pyproject.toml

@@ -59,6 +59,8 @@ dependencies = [
   "omnipath>=1.0.7",
   "pandas>=2.1",
   "pillow>=8",
+  "pooch>=1.6",
+  "pyyaml>=6",
   "scanpy>=1.9.3",
   "scikit-image>=0.25",
   # due to https://github.com/scikit-image/scikit-image/issues/6850 breaks rescale ufunc