diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index ea13a176a..2506c0724 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -2,10 +2,12 @@ from __future__ import unicode_literals import os import re +import io import getpass import posixpath import logging from subprocess import Popen, PIPE +from contextlib import contextmanager, closing try: import pyarrow @@ -139,6 +141,18 @@ class RemoteHDFS(RemoteBASE): with open(to_file, "wb+") as fobj: hdfs.download(from_info.path, fobj) + @contextmanager + def open(self, path_info, mode="r", encoding=None): + assert mode in {"r", "rt", "rb"} + + with self.hdfs(path_info) as hdfs, closing( + hdfs.open(path_info.path, mode="rb") + ) as fd: + if mode == "rb": + yield fd + else: + yield io.TextIOWrapper(fd, encoding=encoding) + def list_cache_paths(self): if not self.exists(self.path_info): return [] diff --git a/dvc/remote/local/__init__.py b/dvc/remote/local/__init__.py index 816cdaa2e..a201ba5bc 100644 --- a/dvc/remote/local/__init__.py +++ b/dvc/remote/local/__init__.py @@ -4,7 +4,7 @@ from copy import copy from dvc.scheme import Schemes from dvc.remote.local.slow_link_detection import slow_link_guard -from dvc.utils.compat import str, fspath_py35 +from dvc.utils.compat import str, fspath_py35, open import os import stat @@ -280,6 +280,10 @@ class RemoteLOCAL(RemoteBASE): name=name, ) + def open(self, path_info, mode="r", encoding=None): + assert mode in {"r", "rt", "rb"} + return open(fspath_py35(path_info), mode=mode, encoding=encoding) + def _group(self, checksum_infos, show_checksums=False): by_md5 = {} diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index 4202ba61a..fee5fd092 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -2,11 +2,13 @@ from __future__ import unicode_literals import errno import itertools +import io import os import getpass import logging import threading from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager, closing from dvc.progress import ProgressCallback from dvc.utils import to_chunks @@ -202,6 +204,18 @@ class RemoteSSH(RemoteBASE): no_progress_bar=no_progress_bar, ) + @contextmanager + def open(self, path_info, mode="r", encoding=None): + assert mode in {"r", "rt", "rb"} + + with self.ssh(path_info) as ssh, closing( + ssh.sftp.file(path_info.path, mode="r") + ) as fd: + if mode == "rb": + yield fd + else: + yield io.TextIOWrapper(fd, encoding=encoding) + def list_cache_paths(self): with self.ssh(self.path_info) as ssh: return list(ssh.walk_files(self.path_info.path)) diff --git a/tests/func/test_api.py b/tests/func/test_api.py index afc71d9ec..adbeebb81 100644 --- a/tests/func/test_api.py +++ b/tests/func/test_api.py @@ -7,6 +7,7 @@ from dvc.main import main from dvc.path_info import URLInfo from dvc.remote.config import RemoteConfig from .test_data_cloud import ( + get_local_url, _should_test_aws, get_aws_url, _should_test_gcp, @@ -23,6 +24,11 @@ from .test_data_cloud import ( # NOTE: staticmethod is only needed in Python 2 +class Local: + should_test = staticmethod(lambda: True) + get_url = staticmethod(get_local_url) + + class S3: should_test = staticmethod(_should_test_aws) get_url = staticmethod(get_aws_url) @@ -53,7 +59,7 @@ class HDFS: get_url = staticmethod(get_hdfs_url) -remote_params = [S3, GCP, Azure, OSS, SSH, HDFS] +remote_params = [Local, S3, GCP, Azure, OSS, SSH, HDFS] @pytest.fixture