Merge pull request #2894 from odd-industries/master

import-url: support directories
This commit is contained in:
Ruslan Kuprieiev
2019-12-06 16:12:04 +02:00
committed by GitHub
2 changed files with 66 additions and 4 deletions
+42 -2
View File
@@ -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:
+24 -2
View File
@@ -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"
) == ""