fs: remove download

dvc/fs/_callback.py

@@ -1,5 +1,5 @@
 from functools import wraps
-from typing import IO, TYPE_CHECKING, Any, Optional, TypeVar, cast
+from typing import IO, TYPE_CHECKING, Any, Dict, Optional, TypeVar, cast
 
 import fsspec
 
@@ -48,8 +48,7 @@ class FsspecCallback(fsspec.Callback):
 
         def make_callback(path1, path2):
             # pylint: disable=assignment-from-none
-            child = self.branch(fs.path.name(path1), path2, {})
-            return self.as_callback(child)
+            return self.branch(fs.path.name(path1), path2, {})
 
         @wraps(fn)
         def func(path1, path2, **kwargs):
@@ -89,6 +88,11 @@ class FsspecCallback(fsspec.Callback):
     ):
         return callback or RichCallback(**rich_kwargs)
 
+    def branch(
+        self, path_1: str, path_2: str, kwargs: Dict[str, Any]
+    ) -> "FsspecCallback":
+        return DEFAULT_CALLBACK
+
 
 class NoOpCallback(FsspecCallback, fsspec.callbacks.NoOpCallback):
     pass

dvc/fs/base.py

@@ -19,7 +19,6 @@ from typing import (
 from funcy import cached_property
 
 from dvc.exceptions import DvcException
-from dvc.utils.fs import as_atomic, makedirs
 from dvc.utils.threadpool import ThreadPoolExecutor
 
 from ._callback import DEFAULT_CALLBACK, FsspecCallback
@@ -360,20 +359,25 @@ class FileSystem:
     def size(self, path: AnyFSPath) -> Optional[int]:
         return self.fs.size(path)
 
-    # pylint: enable=unused-argument
-
-    def download(
+    def get(
         self,
         from_info: AnyFSPath,
         to_info: AnyFSPath,
-        callback: "FsspecCallback" = None,
-        jobs: int = None,
+        callback: "FsspecCallback" = DEFAULT_CALLBACK,
+        batch_size: int = None,
     ):
-        if not self.isdir(from_info):
-            return self.download_file(from_info, to_info, callback=callback)
-
         from .local import localfs
 
+        def get_file(from_info, to_info):
+            localfs.makedirs(localfs.path.parent(to_info), exist_ok=True)
+            branch = callback.branch(from_info, to_info, {})
+            self.get_file(from_info, to_info, callback=branch)
+            callback.relative_update()
+
+        if not self.isdir(from_info):
+            callback.set_size(1)
+            return get_file(from_info, to_info)
+
         pairs = {
             info: localfs.path.join(
                 to_info, *self.path.relparts(info, from_info)
@@ -381,48 +385,24 @@ class FileSystem:
             for info in self.find(from_info)
         }
         if not pairs:
-            makedirs(to_info, exist_ok=True)
+            localfs.makedirs(to_info, exist_ok=True)
             return
 
-        with FsspecCallback.as_tqdm_callback(
-            callback,
-            total=-1,
-            desc=f"Downloading directory {self.path.name(from_info)}",
-            unit="files",
-        ) as cb:
-            cb.set_size(len(pairs))
-            download_files = cb.wrap_and_branch(self.download_file, fs=self)
-            max_workers = jobs or self.jobs
-            # NOTE: unlike pulling/fetching cache, where we need to
-            # download everything we can, not raising an error here might
-            # turn very ugly, as the user might think that he has
-            # downloaded a complete directory, while having a partial one,
-            # which might cause unexpected results in his pipeline.
-            with ThreadPoolExecutor(
-                max_workers=max_workers, cancel_on_error=True
-            ) as executor:
-                list(
-                    executor.imap_unordered(
-                        lambda args: download_files(*args), pairs.items()
-                    )
+        callback.set_size(len(pairs))
+        max_workers = batch_size or self.jobs
+        # NOTE: unlike pulling/fetching cache, where we need to
+        # download everything we can, not raising an error here might
+        # turn very ugly, as the user might think that he has
+        # downloaded a complete directory, while having a partial one,
+        # which might cause unexpected results in his pipeline.
+        with ThreadPoolExecutor(
+            max_workers=max_workers, cancel_on_error=True
+        ) as executor:
+            list(
+                executor.imap_unordered(
+                    lambda args: get_file(*args), pairs.items()
                 )
-
-    def download_file(
-        self,
-        from_info: AnyFSPath,
-        to_info: AnyFSPath,
-        callback: FsspecCallback = None,
-    ):
-        from .local import localfs
-
-        with FsspecCallback.as_tqdm_callback(
-            callback,
-            total=-1,
-            desc=self.path.name(from_info),
-            bytes=True,
-        ) as cb:
-            with as_atomic(localfs, to_info, create_parents=True) as tmp_file:
-                self.get_file(from_info, tmp_file, callback=cb)
+            )
 
 
 class ObjectFileSystem(FileSystem):  # pylint: disable=abstract-method

dvc/fs/utils.py

@@ -3,6 +3,8 @@ import logging
 import os
 from typing import TYPE_CHECKING, List, Optional
 
+from dvc.utils.fs import as_atomic
+
 from ._callback import DEFAULT_CALLBACK, FsspecCallback
 from .base import RemoteActionNotImplemented
 from .local import LocalFileSystem
@@ -45,7 +47,10 @@ def copy(
         return to_fs.put_file(from_path, to_path, callback=callback)
 
     if isinstance(to_fs, LocalFileSystem):
-        return from_fs.download_file(from_path, to_path, callback=callback)
+        from .local import localfs
+
+        with as_atomic(localfs, to_path, create_parents=True) as tmp_file:
+            return from_fs.get_file(from_path, tmp_file, callback=callback)
 
     with from_fs.open(from_path, mode="rb") as fobj:
         size = from_fs.size(from_path)

dvc/output.py

@@ -695,7 +695,14 @@ class Output:
             raise DvcException(msg.format(self.fs_path, name))
 
     def download(self, to, jobs=None):
-        self.fs.download(self.fs_path, to.fs_path, jobs=jobs)
+        from dvc.fs._callback import FsspecCallback
+
+        with FsspecCallback.as_tqdm_callback(
+            total=-1,
+            desc=f"Downloading {self.fs.path.name(self.fs_path)}",
+            unit="files",
+        ) as cb:
+            self.fs.get(self.fs_path, to.fs_path, batch_size=jobs, callback=cb)
 
     def get_obj(self, filter_info=None, **kwargs):
         if self.obj:

dvc/repo/get.py

@@ -21,6 +21,7 @@ def get(url, path, out=None, rev=None, jobs=None):
 
     from dvc.dvcfile import is_valid_filename
     from dvc.external_repo import external_repo
+    from dvc.fs._callback import FsspecCallback
 
     out = resolve_output(path, out)
 
@@ -49,6 +50,14 @@ def get(url, path, out=None, rev=None, jobs=None):
         with external_repo(
             url=url, rev=rev, cache_dir=tmp_dir, cache_types=cache_types
         ) as repo:
-            repo.repo_fs.download(path, os.path.abspath(out), jobs=jobs)
+            fs = repo.repo_fs
+            with FsspecCallback.as_tqdm_callback(
+                total=-1,
+                desc=f"Downloading {fs.path.name(path)}",
+                unit="files",
+            ) as cb:
+                fs.get(
+                    path, os.path.abspath(out), batch_size=jobs, callback=cb
+                )
     finally:
         remove(tmp_dir)

tests/func/test_external_repo.py

@@ -95,7 +95,7 @@ def test_pull_subdir_file(tmp_dir, erepo_dir):
 
     dest = tmp_dir / "file"
     with external_repo(os.fspath(erepo_dir)) as repo:
-        repo.repo_fs.download(
+        repo.repo_fs.get(
             "subdir/file",
             os.fspath(dest),
         )
@@ -191,7 +191,7 @@ def test_subrepos_are_ignored(tmp_dir, erepo_dir):
         subrepo.dvc_gen({"file": "file"}, commit="add files on subrepo")
 
     with external_repo(os.fspath(erepo_dir)) as repo:
-        repo.repo_fs.download(
+        repo.repo_fs.get(
             "dir",
             os.fspath(tmp_dir / "out"),
         )
@@ -237,7 +237,7 @@ def test_subrepos_are_ignored_for_git_tracked_dirs(tmp_dir, erepo_dir):
         subrepo.dvc_gen({"file": "file"}, commit="add files on subrepo")
 
     with external_repo(os.fspath(erepo_dir)) as repo:
-        repo.repo_fs.download(
+        repo.repo_fs.get(
             "dir",
             os.fspath(tmp_dir / "out"),
         )

tests/func/test_fs.py

@@ -3,8 +3,6 @@ from io import BytesIO
 from operator import itemgetter
 from os.path import join
 
-import pytest
-
 from dvc.fs import get_cloud_fs
 from dvc.fs._callback import FsspecCallback
 from dvc.fs.local import LocalFileSystem
@@ -258,33 +256,13 @@ def test_upload_callback(tmp_dir, dvc, cloud):
     assert callback.value == expected_size
 
 
-def test_download_callback(tmp_dir, dvc, cloud, local_cloud):
-    cls, config, _ = get_cloud_fs(dvc, **cloud.config)
-    fs = cls(**config)
-
-    (tmp_dir / "to_upload").write_text("foo")
-    fs.put_file((tmp_dir / "to_upload").fs_path, (cloud / "foo").fs_path)
-    expected_size = fs.size((cloud / "foo").fs_path)
-
-    callback = FsspecCallback()
-    fs.download_file(
-        (cloud / "foo").fs_path,
-        (tmp_dir / "foo").fs_path,
-        callback=callback,
-    )
-
-    assert callback.size == expected_size
-    assert callback.value == expected_size
-    assert (tmp_dir / "foo").read_text() == "foo"
-
-
-def test_download_dir_callback(tmp_dir, dvc, cloud):
+def test_get_dir_callback(tmp_dir, dvc, cloud):
     cls, config, _ = get_cloud_fs(dvc, **cloud.config)
     fs = cls(**config)
     cloud.gen({"dir": {"foo": "foo", "bar": "bar"}})
 
     callback = FsspecCallback()
-    fs.download(
+    fs.get(
         (cloud / "dir").fs_path, (tmp_dir / "dir").fs_path, callback=callback
     )
 
@@ -293,47 +271,14 @@ def test_download_dir_callback(tmp_dir, dvc, cloud):
     assert (tmp_dir / "dir").read_text() == {"foo": "foo", "bar": "bar"}
 
 
-@pytest.mark.parametrize("fs_type", ["git", "dvc"])
-def test_download_callbacks_on_dvc_git_fs(tmp_dir, dvc, scm, fs_type):
-    from dvc.fs.git import GitFileSystem
-
-    gen = tmp_dir.scm_gen if fs_type == "git" else tmp_dir.dvc_gen
-    gen({"dir": {"foo": "foo", "bar": "bar"}, "file": "file"}, commit="gen")
-
-    fs = dvc.dvcfs if fs_type == "dvc" else GitFileSystem(scm=scm, rev="HEAD")
-
-    callback = FsspecCallback()
-    fs.download_file(
-        "file",
-        (tmp_dir / "file2").fs_path,
-        callback=callback,
-    )
-
-    size = os.path.getsize(tmp_dir / "file")
-    assert (tmp_dir / "file2").read_text() == "file"
-    assert callback.size == size
-    assert callback.value == size
-
-    callback = FsspecCallback()
-    fs.download(
-        "dir",
-        (tmp_dir / "dir2").fs_path,
-        callback=callback,
-    )
-
-    assert (tmp_dir / "dir2").read_text() == {"foo": "foo", "bar": "bar"}
-    assert callback.size == 2
-    assert callback.value == 2
-
-
-def test_callback_on_repo_fs(tmp_dir, dvc, scm):
+def test_callback_on_repo_fs(tmp_dir, dvc, scm, mocker):
     tmp_dir.dvc_gen({"dir": {"bar": "bar"}}, commit="dvc")
     tmp_dir.scm_gen({"dir": {"foo": "foo"}}, commit="git")
 
     fs = dvc.repo_fs
 
     callback = FsspecCallback()
-    fs.download(
+    fs.get(
         "dir",
         (tmp_dir / "dir2").fs_path,
         callback=callback,
@@ -344,7 +289,8 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm):
     assert callback.value == 2
 
     callback = FsspecCallback()
-    fs.download(
+    branch = mocker.spy(callback, "branch")
+    fs.get(
         os.path.join("dir", "foo"),
         (tmp_dir / "foo").fs_path,
         callback=callback,
@@ -352,11 +298,18 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm):
 
     size = os.path.getsize(tmp_dir / "dir" / "foo")
     assert (tmp_dir / "foo").read_text() == "foo"
-    assert callback.size == size
-    assert callback.value == size
+    assert callback.size == 1
+    assert callback.value == 1
+
+    assert branch.call_count == 1
+    assert branch.spy_return.size == size
+    assert branch.spy_return.value == size
+
+    branch.reset_mock()
 
     callback = FsspecCallback()
-    fs.download(
+    branch = mocker.spy(callback, "branch")
+    fs.get(
         os.path.join("dir", "bar"),
         (tmp_dir / "bar").fs_path,
         callback=callback,
@@ -364,5 +317,9 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm):
 
     size = os.path.getsize(tmp_dir / "dir" / "bar")
     assert (tmp_dir / "bar").read_text() == "bar"
-    assert callback.size == size
-    assert callback.value == size
+    assert callback.size == 1
+    assert callback.value == 1
+
+    assert branch.call_count == 1
+    assert branch.spy_return.size == size
+    assert branch.spy_return.value == size

tests/func/test_repro.py

@@ -877,10 +877,10 @@ class TestReproAlreadyCached(TestRepro):
         ret = main(["import-url", self.FOO, self.BAR])
         self.assertEqual(ret, 0)
 
-        patch_download = patch.object(
+        patch_get = patch.object(
             LocalFileSystem,
-            "download",
-            side_effect=LocalFileSystem.download,
+            "get",
+            side_effect=LocalFileSystem.get,
             autospec=True,
         )
 
@@ -891,7 +891,7 @@ class TestReproAlreadyCached(TestRepro):
             autospec=True,
         )
 
-        with patch_download as mock_download:
+        with patch_get as mock_download:
             with patch_checkout as mock_checkout:
                 assert main(["unfreeze", "bar.dvc"]) == 0
                 ret = main(["repro", "--force", "bar.dvc"])