mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
Merge pull request #2894 from odd-industries/master
import-url: support directories
This commit is contained in:
+42
-2
@@ -577,11 +577,51 @@ class RemoteBASE(object):
|
||||
if to_info.scheme != "local":
|
||||
raise NotImplementedError
|
||||
|
||||
logger.debug("Downloading '{}' to '{}'".format(from_info, to_info))
|
||||
if self.isdir(from_info):
|
||||
return self._download_dir(
|
||||
from_info, to_info, name, no_progress_bar, file_mode, dir_mode
|
||||
)
|
||||
return self._download_file(
|
||||
from_info, to_info, name, no_progress_bar, file_mode, dir_mode
|
||||
)
|
||||
|
||||
def _download_dir(
|
||||
self, from_info, to_info, name, no_progress_bar, file_mode, dir_mode
|
||||
):
|
||||
file_to_infos = (
|
||||
to_info / file_to_info.relative_to(from_info)
|
||||
for file_to_info in self.walk_files(from_info)
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self.JOBS) as executor:
|
||||
file_from_info = list(self.walk_files(from_info))
|
||||
download_files = partial(
|
||||
self._download_file,
|
||||
name=name,
|
||||
no_progress_bar=True,
|
||||
file_mode=file_mode,
|
||||
dir_mode=dir_mode,
|
||||
)
|
||||
futures = executor.map(
|
||||
download_files, file_from_info, file_to_infos
|
||||
)
|
||||
with Tqdm(
|
||||
futures,
|
||||
total=len(file_from_info),
|
||||
desc="Downloading directory",
|
||||
unit="Files",
|
||||
disable=no_progress_bar,
|
||||
) as futures:
|
||||
return sum(futures)
|
||||
|
||||
def _download_file(
|
||||
self, from_info, to_info, name, no_progress_bar, file_mode, dir_mode
|
||||
):
|
||||
makedirs(to_info.parent, exist_ok=True, mode=dir_mode)
|
||||
|
||||
logger.debug("Downloading '{}' to '{}'".format(from_info, to_info))
|
||||
name = name or to_info.name
|
||||
|
||||
makedirs(to_info.parent, exist_ok=True, mode=dir_mode)
|
||||
tmp_file = tmp_fname(to_info)
|
||||
|
||||
try:
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import pytest
|
||||
|
||||
import os
|
||||
from dvc.remote.s3 import RemoteS3
|
||||
|
||||
from dvc.utils import walk_files
|
||||
from dvc.path_info import PathInfo
|
||||
from tests.remotes import GCP, S3Mocked
|
||||
|
||||
remotes = [GCP, S3Mocked]
|
||||
@@ -132,3 +133,24 @@ def test_isfile(remote):
|
||||
|
||||
for expected, path in test_cases:
|
||||
assert remote.isfile(remote.path_info / path) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("remote", remotes, indirect=True)
|
||||
def test_download_dir(remote, tmpdir):
|
||||
path = str(tmpdir / "data")
|
||||
to_info = PathInfo(path)
|
||||
remote.download(remote.path_info / "data", to_info)
|
||||
assert os.path.isdir(path)
|
||||
data_dir = tmpdir / "data"
|
||||
assert len(list(walk_files(path, None))) == 7
|
||||
assert (data_dir / "alice").read_text(encoding="utf-8") == "alice"
|
||||
assert (data_dir / "alpha").read_text(encoding="utf-8") == "alpha"
|
||||
assert (data_dir / "subdir-file.txt").read_text(
|
||||
encoding="utf-8"
|
||||
) == "subdir"
|
||||
assert (data_dir / "subdir" / "1").read_text(encoding="utf-8") == "1"
|
||||
assert (data_dir / "subdir" / "2").read_text(encoding="utf-8") == "2"
|
||||
assert (data_dir / "subdir" / "3").read_text(encoding="utf-8") == "3"
|
||||
assert (data_dir / "subdir" / "empty_file").read_text(
|
||||
encoding="utf-8"
|
||||
) == ""
|
||||
|
||||
Reference in New Issue
Block a user