Repository: dandi/dandi-cli

Issue:
download: add an option to not download the entire dandiset
the UX problem is that we do allow to download an arbitrary directory, files, etc.

We need an option for a download at dandiset level to not download individual assets, and just get the `dandiset.yaml`.

@satra thinks about `--no-assets` option.

Could be `--download all,dandiset.yaml,assets` so we could make it more flexible in the future if we get other types of entities.

Comments:

Patch:
diff --git a/dandi/cli/cmd_download.py b/dandi/cli/cmd_download.py
index 32b1b481..e44e1c7d 100644
--- a/dandi/cli/cmd_download.py
+++ b/dandi/cli/cmd_download.py
@@ -4,6 +4,29 @@ import click
 from .base import devel_option, map_to_click_exceptions
 
 
+class ChoiceList(click.ParamType):
+    name = "choice-list"
+
+    def __init__(self, values):
+        self.values = set(values)
+
+    def convert(self, value, param, ctx):
+        if value is None or isinstance(value, set):
+            return value
+        selected = set()
+        for v in value.split(","):
+            if v == "all":
+                selected = self.values.copy()
+            elif v in self.values:
+                selected.add(v)
+            else:
+                self.fail(f"{v!r}: invalid value", param, ctx)
+        return selected
+
+    def get_metavar(self, param):
+        return "[" + ",".join(self.values) + ",all]"
+
+
 @click.command()
 @click.option(
     "-o",
@@ -39,6 +62,14 @@ from .base import devel_option, map_to_click_exceptions
     default=6,  # TODO: come up with smart auto-scaling etc
     show_default=True,
 )
+@click.option(
+    "--download",
+    "download_types",
+    type=ChoiceList(["dandiset.yaml", "assets"]),
+    help="Comma-separated list of elements to download",
+    default="all",
+    show_default=True,
+)
 # Might be a cool feature, not unlike verifying a checksum, we verify that
 # downloaded file passes the validator, and if not -- alert
 # @click.option(
@@ -55,16 +86,19 @@ from .base import devel_option, map_to_click_exceptions
 )
 @click.argument("url", nargs=-1)
 @map_to_click_exceptions
-def download(url, output_dir, existing, jobs, format):
+def download(url, output_dir, existing, jobs, format, download_types):
     """Download a file or entire folder from DANDI"""
-    # First boring attempt at click commands being merely an interface to
-    # Python function
-    from ..download import download
+    # We need to import the download module rather than the download function
+    # so that the tests can properly patch the function with a mock.
+    from .. import download
 
-    return download(
+    return download.download(
         url,
         output_dir,
         existing=existing,
         format=format,
-        jobs=jobs,  # develop_debug=develop_debug
+        jobs=jobs,
+        get_metadata="dandiset.yaml" in download_types,
+        get_assets="assets" in download_types,
+        # develop_debug=develop_debug
     )
diff --git a/dandi/download.py b/dandi/download.py
index 28c768d9..92bebded 100644
--- a/dandi/download.py
+++ b/dandi/download.py
@@ -263,7 +263,16 @@ parse_dandi_url = _dandi_url_parser.parse
 follow_redirect = _dandi_url_parser.follow_redirect
 
 
-def download(urls, output_dir, *, format="pyout", existing="error", jobs=1):
+def download(
+    urls,
+    output_dir,
+    *,
+    format="pyout",
+    existing="error",
+    jobs=1,
+    get_metadata=True,
+    get_assets=True,
+):
     # TODO: unduplicate with upload. For now stole from that one
     # We will again use pyout to provide a neat table summarizing our progress
     # with upload etc
@@ -289,7 +298,14 @@ def download(urls, output_dir, *, format="pyout", existing="error", jobs=1):
         # It could handle delegated to generator downloads
         kw["yield_generator_for_fields"] = rec_fields[1:]  # all but path
 
-    gen_ = download_generator(urls, output_dir, existing=existing, **kw)
+    gen_ = download_generator(
+        urls,
+        output_dir,
+        existing=existing,
+        get_metadata=get_metadata,
+        get_assets=get_assets,
+        **kw,
+    )
 
     # TODOs:
     #  - redo frontends similarly to how command_ls did it
@@ -316,6 +332,8 @@ def download_generator(
     assets_it=None,
     yield_generator_for_fields=None,
     existing="error",
+    get_metadata=True,
+    get_assets=True,
 ):
     """A generator for downloads of files, folders, or entire dandiset from DANDI
     (as identified by URL)
@@ -382,7 +400,7 @@ def download_generator(
         # more efficient download if files are just renamed etc
 
         # Handle our so special dandiset.yaml
-        if dandiset:
+        if dandiset and get_metadata:
             for resp in _populate_dandiset_yaml(
                 dandiset_path,
                 dandiset.get("metadata", {}).get("dandiset", {}),
@@ -390,6 +408,9 @@ def download_generator(
             ):
                 yield dict(path=dandiset_metadata_file, **resp)
 
+        if not get_assets:
+            return
+
         for asset in assets:
             # unavoidable ugliness since girder and API have different "scopes" for
             # identifying an asset


Test patch:
diff --git a/dandi/cli/tests/test_download.py b/dandi/cli/tests/test_download.py
new file mode 100644
index 00000000..697ba9a1
--- /dev/null
+++ b/dandi/cli/tests/test_download.py
@@ -0,0 +1,73 @@
+import os
+import click
+from click.testing import CliRunner
+from ..command import download
+
+
+def test_download_defaults(mocker):
+    mock_download = mocker.patch("dandi.download.download")
+    r = CliRunner().invoke(download)
+    assert r.exit_code == 0
+    mock_download.assert_called_once_with(
+        (),
+        os.curdir,
+        existing="refresh",
+        format="pyout",
+        jobs=6,
+        get_metadata=True,
+        get_assets=True,
+    )
+
+
+def test_download_all_types(mocker):
+    mock_download = mocker.patch("dandi.download.download")
+    r = CliRunner().invoke(download, ["--download", "all"])
+    assert r.exit_code == 0
+    mock_download.assert_called_once_with(
+        (),
+        os.curdir,
+        existing="refresh",
+        format="pyout",
+        jobs=6,
+        get_metadata=True,
+        get_assets=True,
+    )
+
+
+def test_download_metadata_only(mocker):
+    mock_download = mocker.patch("dandi.download.download")
+    r = CliRunner().invoke(download, ["--download", "dandiset.yaml"])
+    assert r.exit_code == 0
+    mock_download.assert_called_once_with(
+        (),
+        os.curdir,
+        existing="refresh",
+        format="pyout",
+        jobs=6,
+        get_metadata=True,
+        get_assets=False,
+    )
+
+
+def test_download_assets_only(mocker):
+    mock_download = mocker.patch("dandi.download.download")
+    r = CliRunner().invoke(download, ["--download", "assets"])
+    assert r.exit_code == 0
+    mock_download.assert_called_once_with(
+        (),
+        os.curdir,
+        existing="refresh",
+        format="pyout",
+        jobs=6,
+        get_metadata=False,
+        get_assets=True,
+    )
+
+
+def test_download_bad_type(mocker):
+    mock_download = mocker.patch("dandi.download.download")
+    r = CliRunner().invoke(download, ["--download", "foo"], standalone_mode=False)
+    assert r.exit_code != 0
+    assert isinstance(r.exception, click.UsageError)
+    assert str(r.exception) == "'foo': invalid value"
+    mock_download.assert_not_called()
diff --git a/dandi/tests/test_download.py b/dandi/tests/test_download.py
index 2d87607c..9bf65193 100644
--- a/dandi/tests/test_download.py
+++ b/dandi/tests/test_download.py
@@ -220,6 +220,38 @@ def test_download_000027(url, tmpdir):
     download(url, tmpdir, existing="refresh")  # TODO: check that skipped (the same)
 
 
+@pytest.mark.parametrize(
+    "url",
+    [  # Should go through API
+        "https://dandiarchive.org/dandiset/000027/0.200721.2222",
+        # Drafts do not go through API ATM, but that should not be visible to user
+        "https://dandiarchive.org/dandiset/000027/draft",
+    ],
+)
+def test_download_000027_metadata_only(url, tmpdir):
+    ret = download(url, tmpdir, get_assets=False)
+    assert not ret  # we return nothing ATM, might want to "generate"
+    dsdir = tmpdir / "000027"
+    downloads = (x.relto(dsdir) for x in dsdir.visit())
+    assert sorted(downloads) == ["dandiset.yaml"]
+
+
+@pytest.mark.parametrize(
+    "url",
+    [  # Should go through API
+        "https://dandiarchive.org/dandiset/000027/0.200721.2222",
+        # Drafts do not go through API ATM, but that should not be visible to user
+        "https://dandiarchive.org/dandiset/000027/draft",
+    ],
+)
+def test_download_000027_assets_only(url, tmpdir):
+    ret = download(url, tmpdir, get_metadata=False)
+    assert not ret  # we return nothing ATM, might want to "generate"
+    dsdir = tmpdir / "000027"
+    downloads = (x.relto(dsdir) for x in dsdir.visit())
+    assert sorted(downloads) == ["sub-RAT123", op.join("sub-RAT123", "sub-RAT123.nwb")]
+
+
 def test_girder_tqdm(monkeypatch):
     # smoke test to ensure we do not blow up
     def raise_assertion_error(*args, **kwargs):