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
18 changes: 18 additions & 0 deletions tests/unit/remote/test_remote_dir.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# -*- coding: utf-8 -*-
import pytest
import os

from dvc.remote.s3 import RemoteS3

Copy link
Contributor

@efiop efiop Dec 6, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The reason why tests are failing 🙂

Suggested change
from dvc.utils import walk_files

Expand Down Expand Up @@ -132,3 +133,20 @@ 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 = os.fspath(tmpdir / "data")
efiop marked this conversation as resolved.
Show resolved Hide resolved
to_info = os.PathInfo(path)
efiop marked this conversation as resolved.
Show resolved Hide resolved
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") == ""