Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

import-url: support directories #2894

Merged
merged 14 commits into from
Dec 6, 2019
44 changes: 42 additions & 2 deletions dvc/remote/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -547,11 +547,51 @@ def download(
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
)

name = name or to_info.name
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:
efiop marked this conversation as resolved.
Show resolved Hide resolved
file_from_info = list(self.walk_files(from_info))
efiop marked this conversation as resolved.
Show resolved Hide resolved
download_files = partial(
self._download_file,
name=name,
no_progress_bar=True,
efiop marked this conversation as resolved.
Show resolved Hide resolved
file_mode=file_mode,
dir_mode=dir_mode,
)
futures = executor.map(
download_files, file_from_info, file_to_infos
)
with Tqdm(
efiop marked this conversation as resolved.
Show resolved Hide resolved
futures,
total=len(file_from_info),
desc="Downloading directory",
efiop marked this conversation as resolved.
Show resolved Hide resolved
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

tmp_file = tmp_fname(to_info)

try:
Expand Down
26 changes: 24 additions & 2 deletions tests/unit/remote/test_remote_dir.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
# -*- coding: utf-8 -*-
import pytest

import os
from dvc.remote.s3 import RemoteS3

from dvc.utils import walk_files
efiop marked this conversation as resolved.
Show resolved Hide resolved
from dvc.path_info import PathInfo
from tests.remotes import GCP, S3Mocked

remotes = [GCP, S3Mocked]
Expand Down Expand Up @@ -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"
) == ""