Support for partial download of zarrs

dandi/dandiapi.py

@@ -1240,16 +1240,29 @@ def get_asset_by_path(self, path: str) -> RemoteAsset:
         paths (e.g., ``../``) within it.
         """
         path = self._normalize_path(path)
-        try:
-            # Weed out any assets that happen to have the given path as a
-            # proper prefix:
-            (asset,) = (
-                a for a in self.get_assets_with_path_prefix(path) if a.path == path
-            )
-        except ValueError:
-            raise NotFoundError(f"No asset at path {path!r}")
-        else:
-            return asset
+        # Weed out any assets that happen to have the given path as a
+        # proper prefix:
+        assets = [a for a in self.get_assets_with_path_prefix(path) if a.path == path]
+        if assets:
+            if len(assets) > 1:
+                lgr.warning(
+                    "Multiple assets found at path %r; returning the first one",
+                    path,
+                )
+            asset = assets[0]
+        elif not assets:
+            zarr_suf = ".zarr/"
+            if (zarr_suf in path) and (zarr_suffix_idx := path.index(zarr_suf)):
+                # If path ends with .zarr/, we might have a zarr asset without
+                # a trailing slash in the path.  Try again:
+                zarr_path_len = zarr_suffix_idx + len(zarr_suf)
+                zarr_path = path[: zarr_path_len - 1]  # -1 for trailing /
+                subpath = path[len(zarr_path) :]
+                asset = self.get_asset_by_path(zarr_path)
+                asset.subpath = subpath
+            else:
+                raise NotFoundError(f"No asset at path {path!r}")
+        return asset
 
     def download_directory(
         self,
@@ -1378,6 +1391,8 @@ class BaseRemoteAsset(ABC, APIBase):
     identifier: str = Field(alias="asset_id")
     #: The asset's (forward-slash-separated) path
     path: str
+    #: The path within the asset, as e.g. path in a zarr/
+    subpath: str | None = Field(default=None)
     #: The size of the asset in bytes
     size: int
     #: The date at which the asset was created

dandi/download.py

@@ -226,6 +226,7 @@ class Downloader:
     url: ParsedDandiURL
     output_dir: InitVar[str | Path]
     output_prefix: Path = field(init=False)
+    #: just a convenience combination of output_dir and output_prefix
     output_path: Path = field(init=False)
     existing: DownloadExisting
     get_metadata: bool
@@ -333,6 +334,12 @@ def download_generator(self) -> Iterator[dict]:
                             asset.path,
                         )
                         mtime = asset.modified
+                    if asset.subpath:
+                        lgr.warning(
+                            "No downloading of subpaths within blobs yet. Got %s for %s",
+                            asset.subpath,
+                            asset.path,
+                        )
                     _download_generator = _download_file(
                         asset.get_download_file_iter(),
                         download_path,
@@ -352,7 +359,8 @@ def download_generator(self) -> Iterator[dict]:
                     ), f"Asset {asset.path} is neither blob nor Zarr"
                     _download_generator = _download_zarr(
                         asset,
-                        download_path,
+                        prefix=asset.subpath,
+                        download_path=download_path,
                         toplevel_path=self.output_path,
                         existing=self.existing,
                         jobs=self.jobs_per_zarr,
@@ -812,6 +820,7 @@ def _download_file(
         lgr.warning("downloader logic: We should not be here!")
 
     final_digest = None
+
     if downloaded_digest and not resuming:
         assert downloaded_digest is not None
         final_digest = downloaded_digest.hexdigest()  # we care only about hex
@@ -977,6 +986,7 @@ def _download_zarr(
     toplevel_path: str | Path,
     existing: DownloadExisting,
     lock: Lock,
+    prefix: str | None = None,
     jobs: int | None = None,
 ) -> Iterator[dict]:
     # Avoid heavy import by importing within function:
@@ -993,7 +1003,7 @@ def digest_callback(path: str, algoname: str, d: str) -> None:
             digests[path] = d
 
     def downloads_gen():
-        for entry in asset.iterfiles():
+        for entry in asset.iterfiles(prefix=prefix):
             entries.append(entry)
             etag = entry.digest
             assert etag.algorithm is DigestType.md5