fs: prefer generic.copy over fsspec fs.get()

This commit is contained in:
Peter Rowlands
2023-07-24 17:01:34 +09:00
committed by Peter Rowlands (변기호)
parent a9d0ac2ba9
commit 7806afd4ad
5 changed files with 44 additions and 19 deletions
+1 -1
View File
@@ -50,7 +50,7 @@ class Dependency(Output):
self.fs_path = self.fs.path.version_path(self.fs_path, self.meta.version_id)
def download(self, to, jobs=None):
fs_download(self.fs, self.fs_path, to, jobs=jobs)
fs_download(self.fs, self.fs_path, to.fs_path, jobs=jobs)
def save(self):
super().save()
+31 -3
View File
@@ -1,3 +1,4 @@
from typing import Optional
from urllib.parse import urlparse
from dvc_http import HTTPFileSystem, HTTPSFileSystem # noqa: F401
@@ -18,7 +19,7 @@ from dvc_objects.fs import ( # noqa: F401
system,
utils,
)
from dvc_objects.fs.base import AnyFSPath, FileSystem # noqa: F401
from dvc_objects.fs.base import AnyFSPath, FileSystem # noqa: F401, TCH001
from dvc_objects.fs.errors import ( # noqa: F401
AuthError,
ConfigError,
@@ -48,12 +49,39 @@ known_implementations.update(
# pylint: enable=unused-import
def download(fs, fs_path, to, jobs=None):
def download(fs: "FileSystem", fs_path: str, to: str, jobs: Optional[int] = None):
with Callback.as_tqdm_callback(
desc=f"Downloading {fs.path.name(fs_path)}",
unit="files",
) as cb:
fs.get(fs_path, to.fs_path, batch_size=jobs, callback=cb)
# NOTE: We use dvc-objects generic.copy over fs.get since it makes file
# download atomic and avoids fsspec glob/regex path expansion.
if fs.isdir(fs_path):
from_infos = [
path
for path in fs.find(fs_path)
if not path.endswith(fs.path.flavour.sep)
]
if not from_infos:
return localfs.makedirs(to, exist_ok=True)
to_infos = [
localfs.path.join(to, *fs.path.relparts(info, fs_path))
for info in from_infos
]
else:
from_infos = [fs_path]
to_infos = [to]
cb.set_size(len(from_infos))
jobs = jobs or fs.jobs
generic.copy(
fs,
from_infos,
localfs,
to_infos,
callback=cb,
batch_size=jobs,
)
def parse_external_url(url, config=None):
+7 -12
View File
@@ -22,7 +22,6 @@ class GetDVCFileError(DvcException):
def get(url, path, out=None, rev=None, jobs=None, force=False, config=None):
from dvc.config import Config
from dvc.dvcfile import is_valid_filename
from dvc.fs.callbacks import Callback
from dvc.repo import Repo
out = resolve_output(path, out, force=force)
@@ -40,6 +39,7 @@ def get(url, path, out=None, rev=None, jobs=None, force=False, config=None):
uninitialized=True,
config=config,
) as repo:
from dvc.fs import download
from dvc.fs.data import DataFileSystem
fs: Union[DataFileSystem, "DVCFileSystem"]
@@ -49,14 +49,9 @@ def get(url, path, out=None, rev=None, jobs=None, force=False, config=None):
else:
fs = repo.dvcfs
fs_path = fs.from_os_path(path)
with Callback.as_tqdm_callback(
desc=f"Downloading {fs.path.name(path)}",
unit="files",
) as cb:
fs.get(
fs_path,
os.path.abspath(out),
batch_size=jobs,
callback=cb,
)
download(
fs,
fs_path,
os.path.abspath(out),
jobs=jobs,
)
+1 -1
View File
@@ -14,4 +14,4 @@ def get_url(url, out=None, *, config=None, jobs=None, force=False):
src_fs, src_path = parse_external_url(url, config)
if not src_fs.exists(src_path):
raise URLMissingError(url)
download(src_fs, src_path, out, jobs=jobs)
download(src_fs, src_path, out.fs_path, jobs=jobs)
+4 -2
View File
@@ -10,7 +10,7 @@ from funcy import lsplit
from dvc.cli import main
from dvc.dvcfile import LOCK_FILE, PROJECT_FILE
from dvc.exceptions import CyclicGraphError, ReproductionError
from dvc.fs import LocalFileSystem, system
from dvc.fs import system
from dvc.output import Output
from dvc.stage import PipelineStage, Stage
from dvc.stage.cache import RunCacheNotSupported
@@ -624,12 +624,14 @@ class TestReproAlreadyCached:
assert stage.outs[0].hash_info != saved_stage.outs[0].hash_info
def test_force_import(self, mocker, tmp_dir, dvc):
from dvc.dependency import base
tmp_dir.dvc_gen("foo", "foo")
ret = main(["import-url", "foo", "bar"])
assert ret == 0
spy_get = mocker.spy(LocalFileSystem, "get")
spy_get = mocker.spy(base, "fs_download")
spy_checkout = mocker.spy(Output, "checkout")
assert main(["unfreeze", "bar.dvc"]) == 0