Fix partial Zarr download by URL: per-URL filters, missing subpaths, trailing slashes

.github/workflows/run-tests.yml

@@ -60,7 +60,9 @@ jobs:
             python: 3.14
             mode: obolibrary-only
           - os: ubuntu-latest
-            python: '3.11'
+            # cannot use dev-deps on 3.11 now, see
+            # https://github.com/hdmf-dev/hdmf-zarr/issues/406
+            python: '3.12'
             mode: dev-deps
           - os: ubuntu-latest
             python: 3.14

dandi/dandiarchive.py

@@ -54,9 +54,13 @@
 from .dandiapi import BaseRemoteAsset, DandiAPIClient, RemoteDandiset
 from .exceptions import FailedToConnectError, NotFoundError, UnknownURLError
 from .utils import get_instance, get_retry_after
+from .zarr_filter import ZarrFilter
 
 lgr = get_logger()
 
+#: `ZARR_EXTENSIONS` in the form `str.endswith()` takes
+_ZARR_SUFFIXES = tuple(ZARR_EXTENSIONS)
+
 
 @dataclass
 class ParsedDandiURL(ABC):
@@ -201,6 +205,20 @@ def get_asset_download_path(
         """
         ...
 
+    def get_zarr_filter(self) -> list[ZarrFilter]:
+        """
+        Returns the filters restricting which entries within the Zarr assets
+        returned by `get_assets()` should be downloaded.  An empty list means
+        that no restriction is implied by the URL and all entries are to be
+        downloaded.
+
+        Only `AssetZarrEntryURL` — a URL pointing inside a Zarr asset — returns
+        a non-empty list.
+
+        :meta private:
+        """
+        return []
+
     @abstractmethod
     def is_under_download_path(self, path: str) -> bool:
         """
@@ -474,13 +492,38 @@ def split_zarr_location(location: str) -> tuple[str, str] | None:
     """
     parts = [p for p in location.split("/") if p]
     for i, part in enumerate(parts):
-        if any(part.endswith(ext) for ext in ZARR_EXTENSIONS):
+        if part.endswith(_ZARR_SUFFIXES):
             asset_path = "/".join(parts[: i + 1])
             zarr_subpath = "/".join(parts[i + 1 :])
             return (asset_path, zarr_subpath) if zarr_subpath else None
     return None
 
 
+def at_zarr_boundary(location: str) -> bool:
+    """Whether ``location``'s last component ends with a Zarr extension.
+
+    Parameters
+    ----------
+    location : str
+        A POSIX-style path, e.g. ``"sub-1/file.ome.zarr/"``.
+
+    Returns
+    -------
+    bool
+        True if the path ends at a Zarr asset (ignoring a trailing slash).
+
+    Examples
+    --------
+    >>> at_zarr_boundary("sub-1/file.ome.zarr/")
+    True
+    >>> at_zarr_boundary("sub-1/file.ome.zarr/0/0")  # below the boundary
+    False
+    >>> at_zarr_boundary("sub-1/")
+    False
+    """
+    return location.rstrip("/").endswith(_ZARR_SUFFIXES)
+
+
 @dataclass
 class AssetZarrEntryURL(SingleAssetURL):
     """Parsed from a URL that points into entries within a Zarr asset.
@@ -489,8 +532,10 @@ class AssetZarrEntryURL(SingleAssetURL):
     produce ``asset_path="sub-1/file.ome.zarr"`` and ``zarr_subpath="0/0/0"``.
     """
 
-    asset_path: str  # e.g., "sub-1/file.ome.zarr"
-    zarr_subpath: str  # e.g., "0/0/0"
+    #: The path of the Zarr asset, e.g. ``"sub-1/file.ome.zarr"``
+    asset_path: str
+    #: The path within the Zarr asset, e.g. ``"0/0/0"``
+    zarr_subpath: str
 
     def get_assets(
         self, client: DandiAPIClient, order: str | None = None, strict: bool = False
@@ -512,6 +557,17 @@ def get_assets(
         with _maybe_strict(strict):
             yield dandiset.get_asset_by_path(self.asset_path)
 
+    def get_zarr_filter(self) -> list[ZarrFilter]:
+        """Restrict the download to the entries at or under `zarr_subpath`.
+
+        :meta private:
+        """
+        if not self.zarr_subpath:
+            # `parse_dandi_url()` never produces this, but the class is public
+            # and an empty subpath would otherwise reject every entry.
+            return []
+        return [ZarrFilter("path", self.zarr_subpath)]
+
 
 @dataclass
 class AssetFolderURL(MultiAssetURL):
@@ -905,6 +961,28 @@ def parse(
                     version_id=version_id,
                     path=location,
                 )
+            elif (zarr_split := split_zarr_location(location)) is not None:
+                # The location crosses a zarr boundary.  This is checked
+                # before the folder case, as a path within a zarr never names
+                # a folder of assets: entries within a zarr are not assets.
+                asset_path, zarr_subpath = zarr_split
+                parsed_url = AssetZarrEntryURL(
+                    instance=instance,
+                    dandiset_id=dandiset_id,
+                    version_id=version_id,
+                    asset_path=asset_path,
+                    zarr_subpath=zarr_subpath,
+                )
+            elif location.endswith("/") and at_zarr_boundary(location):
+                # `.../x.zarr/` names the zarr asset itself; a folder of
+                # assets by that name could never hold it, as the asset's own
+                # path does not end in a slash.
+                parsed_url = AssetItemURL(
+                    instance=instance,
+                    dandiset_id=dandiset_id,
+                    version_id=version_id,
+                    path=location.rstrip("/"),
+                )
             elif location.endswith("/"):
                 parsed_url = AssetFolderURL(
                     instance=instance,
@@ -913,24 +991,12 @@ def parse(
                     path=location,
                 )
             else:
-                # Check if location crosses a zarr boundary
-                zarr_split = split_zarr_location(location)
-                if zarr_split is not None:
-                    asset_path, zarr_subpath = zarr_split
-                    parsed_url = AssetZarrEntryURL(
-                        instance=instance,
-                        dandiset_id=dandiset_id,
-                        version_id=version_id,
-                        asset_path=asset_path,
-                        zarr_subpath=zarr_subpath,
-                    )
-                else:
-                    parsed_url = AssetItemURL(
-                        instance=instance,
-                        dandiset_id=dandiset_id,
-                        version_id=version_id,
-                        path=location,
-                    )
+                parsed_url = AssetItemURL(
+                    instance=instance,
+                    dandiset_id=dandiset_id,
+                    version_id=version_id,
+                    path=location,
+                )
         elif asset_id:
             if dandiset_id is None:
                 parsed_url = BaseAssetIDURL(instance=instance, asset_id=asset_id)

dandi/delete.py

@@ -142,6 +142,13 @@ def register_url(self, url: str) -> None:
             assert parsed_url.dandiset_id is not None
             self.register_dandiset(parsed_url.instance, parsed_url.dandiset_id)
         else:
+            if parsed_url.get_zarr_filter():
+                # The URL points inside a Zarr asset, but `get_assets()` yields
+                # the whole asset, so deleting it would take the entire Zarr.
+                raise NotImplementedError(
+                    "Cannot delete individual entries within a Zarr asset;"
+                    f" {url} points inside one"
+                )
             if parsed_url.version_id is None:
                 parsed_url.version_id = DRAFT
             self.register_assets_url(url, parsed_url)

dandi/download.py

@@ -48,7 +48,6 @@
 from .dandiapi import AssetType, BaseRemoteZarrAsset, RemoteDandiset
 from .dandiarchive import (
     AssetItemURL,
-    AssetZarrEntryURL,
     DandisetURL,
     ParsedDandiURL,
     SingleAssetURL,
@@ -122,23 +121,20 @@ def download(
 
     parsed_urls = [parse_dandi_url(u, glob=path_type is PathType.GLOB) for u in urls]
 
-    # Parse zarr entry filters
-    zarr_entry_filter: Callable[[str], bool] | None = None
-    all_zf: list[ZarrFilter] = []
-    if zarr_filters:
-        for spec in zarr_filters:
-            all_zf.extend(parse_zarr_filter(spec))
+    # Parse the explicit ``--zarr`` filters.  Filters implied by a URL that
+    # points inside a Zarr asset are resolved per-URL by each `Downloader`, so
+    # that the subpath of one URL does not restrict the download of another.
+    explicit_zarr_filters: list[ZarrFilter] = []
+    for spec in zarr_filters:
+        explicit_zarr_filters.extend(parse_zarr_filter(spec))
 
-    # Merge URL-derived path filters from AssetZarrEntryURL
-    for purl in parsed_urls:
-        if isinstance(purl, AssetZarrEntryURL):
-            all_zf.append(ZarrFilter("path", purl.zarr_subpath))
-
-    if all_zf:
-        zarr_entry_filter = make_zarr_entry_filter(all_zf)
-
-    if sync and zarr_entry_filter is not None:
-        raise ValueError("--sync and --zarr cannot be used together")
+    if sync:
+        if explicit_zarr_filters:
+            raise ValueError("--sync and --zarr cannot be used together")
+        if any(purl.get_zarr_filter() for purl in parsed_urls):
+            raise ValueError(
+                "--sync cannot be used with a URL pointing inside a Zarr asset"
+            )
 
     # dandi.cli.formatters are used in cmd_ls to provide switchable
     pyout_style = pyouts.get_style(hide_if_missing=False)
@@ -176,7 +172,7 @@ def download(
             preserve_tree=preserve_tree,
             jobs_per_zarr=jobs_per_zarr,
             on_error="yield" if format is DownloadFormat.PYOUT else "raise",
-            zarr_entry_filter=zarr_entry_filter,
+            zarr_filters=explicit_zarr_filters,
             **kw,
         )
         for purl in parsed_urls
@@ -272,7 +268,17 @@ class Downloader:
     preserve_tree: bool
     jobs_per_zarr: int | None
     on_error: Literal["raise", "yield"]
-    zarr_entry_filter: Callable[[str], bool] | None = None
+    #: Filters from the ``--zarr`` option; they apply to every Zarr asset, and
+    #: matching no entries in a given asset is not an error
+    zarr_filters: list[ZarrFilter] = field(default_factory=list)
+    #: Filters implied by `url` pointing inside a Zarr asset.  Unlike
+    #: `zarr_filters`, these name entries the user explicitly asked for, so
+    #: matching none of them is an error
+    url_zarr_filters: list[ZarrFilter] = field(init=False, default_factory=list)
+    #: Predicate for selecting entries within a Zarr asset, combining
+    #: `zarr_filters` with `url_zarr_filters`; `None` means that every entry is
+    #: to be downloaded
+    zarr_entry_filter: Callable[[str], bool] | None = field(init=False, default=None)
     #: which will be set .gen to assets.  Purpose is to make it possible to get
     #: summary statistics while already downloading.  TODO: reimplement
     #: properly!
@@ -292,6 +298,11 @@ def __post_init__(self, output_dir: str | Path) -> None:
         else:
             self.output_prefix = Path()
         self.output_path = Path(output_dir, self.output_prefix)
+        self.url_zarr_filters = self.url.get_zarr_filter()
+        all_filters = self.url_zarr_filters + list(self.zarr_filters)
+        self.zarr_entry_filter = (
+            make_zarr_entry_filter(all_filters) if all_filters else None
+        )
 
     def is_dandiset_yaml(self) -> bool:
         return isinstance(self.url, AssetItemURL) and self.url.path == "dandiset.yaml"
@@ -345,6 +356,17 @@ def download_generator(self) -> Iterator[dict]:
                 download_path = Path(self.output_path, path)
                 path = str(self.output_prefix / path)
 
+                if self.url_zarr_filters and asset.asset_type is not AssetType.ZARR:
+                    # The URL named a path inside the asset, but the asset is
+                    # not a Zarr, so there is nothing to descend into.
+                    yield {
+                        "path": path,
+                        "status": "error",
+                        "message": f"Asset {asset.path!r} is not a Zarr asset,"
+                        " so it has no entries to download",
+                    }
+                    continue
+
                 try:
                     metadata = asset.get_raw_metadata()
                 except NotFoundError as e:
@@ -398,6 +420,7 @@ def download_generator(self) -> Iterator[dict]:
                         jobs=self.jobs_per_zarr,
                         lock=lock,
                         zarr_entry_filter=self.zarr_entry_filter,
+                        required_filters=self.url_zarr_filters,
                     )
 
                 def _progress_filter(gen):
@@ -1041,6 +1064,7 @@ def _download_zarr(
     lock: Lock,
     jobs: int | None = None,
     zarr_entry_filter: Callable[[str], bool] | None = None,
+    required_filters: Sequence[ZarrFilter] = (),
 ) -> Iterator[dict]:
     # Avoid heavy import by importing within function:
     from .support.digests import get_zarr_checksum
@@ -1050,16 +1074,33 @@ def _download_zarr(
     entries: list = []
     digests: dict[str, str] = {}
     pc = ProgressCombiner(zarr_size=asset.size)
+    # `zarr_entry_filter` is the OR of the required filters and the ``--zarr``
+    # ones, so a non-empty `entries` does not mean the required filters
+    # matched; track them separately.
+    required_match = (
+        make_zarr_entry_filter(list(required_filters)) if required_filters else None
+    )
+    matched_required = False
+
+    def unmatched_required_error() -> dict:
+        patterns = ", ".join(repr(f.pattern) for f in required_filters)
+        return {
+            "status": "error",
+            "message": f"No entries in the Zarr asset match {patterns}",
+        }
 
     def digest_callback(path: str, algoname: str, d: str) -> None:
         if algoname == "md5":
             digests[path] = d
 
     def downloads_gen():
+        nonlocal matched_required
         for entry in asset.iterfiles():
             entry_path = str(entry)
             if zarr_entry_filter is not None and not zarr_entry_filter(entry_path):
                 continue
+            if required_match is not None and required_match(entry_path):
+                matched_required = True
             entries.append(entry)
             etag = entry.digest
             assert etag.algorithm is DigestType.md5
@@ -1094,11 +1135,19 @@ def downloads_gen():
             if final_out is not None:
                 break
         else:
-            if zarr_entry_filter is not None:
+            if required_filters and not matched_required:
+                yield unmatched_required_error()
+            elif zarr_entry_filter is not None:
                 # Filter matched no entries; still report completion
                 yield {"status": "done"}
             return
 
+    if required_filters and not matched_required:
+        # Entries downloaded for the ``--zarr`` filters do not make up for the
+        # ones the URL named but the asset does not have.
+        yield unmatched_required_error()
+        return
+
     if zarr_entry_filter is not None:
         # Partial download: skip deleting extra local files and skip
         # whole-zarr checksum verification (individual file checksums