Merge pull request #2393 from Suor/stream-more

api: stream api.open() for local, SSH and HDFS

dvc/remote/gs.py

@@ -3,6 +3,7 @@ from __future__ import unicode_literals
 import logging
 from datetime import timedelta
 from funcy import cached_property
+from dvc.utils.compat import FileNotFoundError
 
 try:
     from google.cloud import storage
@@ -15,7 +16,6 @@ from dvc.exceptions import DvcException
 from dvc.path_info import CloudURLInfo
 from dvc.scheme import Schemes
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -104,4 +104,6 @@ class RemoteGS(RemoteBASE):
 
         bucket = self.gs.bucket(path_info.bucket)
         blob = bucket.get_blob(path_info.path)
+        if blob is None:
+            raise FileNotFoundError
         return blob.generate_signed_url(expiration=expiration)

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
@@ -15,7 +17,7 @@ except ImportError:
 from dvc.config import Config
 from dvc.scheme import Schemes
 
-from dvc.utils.compat import urlparse
+from dvc.utils.compat import urlparse, FileNotFoundError
 from dvc.utils import fix_env, tmp_fname
 
 from .pool import get_connection
@@ -139,6 +141,25 @@ 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"}
+
+        try:
+            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)
+        except IOError as e:
+            # Empty .errno and not specific enough error class in pyarrow,
+            # see https://issues.apache.org/jira/browse/ARROW-6248
+            if "file does not exist" in str(e):
+                raise FileNotFoundError(*e.args)
+            raise
+
     def list_cache_paths(self):
         if not self.exists(self.path_info):
             return []

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 = {}
 

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))

dvc/repo/__init__.py

@@ -2,7 +2,7 @@ from __future__ import unicode_literals
 
 import os
 import logging
-
+from contextlib import contextmanager
 from itertools import chain
 
 from funcy import cached_property
@@ -17,7 +17,7 @@ from dvc.exceptions import (
 from dvc.ignore import DvcIgnoreFilter
 from dvc.path_info import PathInfo
 from dvc.remote.base import RemoteActionNotImplemented
-from dvc.utils.compat import open as _open, fspath_py35
+from dvc.utils.compat import open as _open, fspath_py35, FileNotFoundError
 from dvc.utils import relpath
 
 logger = logging.getLogger(__name__)
@@ -456,8 +456,16 @@ class Repo(object):
         path_parts = os.path.normpath(path).split(os.path.sep)
         return self.DVC_DIR in path_parts
 
+    @contextmanager
     def open(self, path, remote=None, mode="r", encoding=None):
         """Opens a specified resource as a file descriptor"""
+        try:
+            with self._open(path, remote, mode, encoding) as fd:
+                yield fd
+        except FileNotFoundError:
+            raise OutputFileMissingError(relpath(path, self.root_dir))
+
+    def _open(self, path, remote=None, mode="r", encoding=None):
         out, = self.find_outs_by_path(path)
         if out.isdir():
             raise ValueError("Can't open a dir")
@@ -477,10 +485,6 @@ class Repo(object):
                 cache_info = out.get_used_cache(remote=remote)
                 self.cloud.pull(cache_info, remote=remote)
 
-            # Since pull may just skip with a warning, we need to check it here
-            if not os.path.exists(cache_file):
-                raise OutputFileMissingError(relpath(path, self.root_dir))
-
             return _open(cache_file, mode=mode, encoding=encoding)
 
     @cached_property

dvc/utils/compat.py

@@ -119,6 +119,7 @@ if is_py2:
     cast_bytes_py2 = cast_bytes
     makedirs = _makedirs
     range = xrange  # noqa: F821
+    FileNotFoundError = IOError
 
     import StringIO
     import io
@@ -165,6 +166,7 @@ elif is_py3:
     open = open  # noqa: F821
     cast_bytes_py2 = no_code
     range = range  # noqa: F821
+    FileNotFoundError = FileNotFoundError
 
 
 # Backport os.fspath() from Python 3.6

dvc/utils/http.py

@@ -3,6 +3,8 @@ import io
 
 import requests
 
+from dvc.utils.compat import FileNotFoundError
+
 
 @contextmanager
 def open_url(url, mode="r", encoding=None):
@@ -30,6 +32,8 @@ def iter_url(url, chunk_size=io.DEFAULT_BUFFER_SIZE):
     def request(headers=None):
         the_url = url() if callable(url) else url
         response = requests.get(the_url, stream=True, headers=headers)
+        if response.status_code == 404:
+            raise FileNotFoundError("Can't open {}".format(the_url))
         response.raise_for_status()
         return response
 

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)
@@ -54,6 +60,7 @@ class HDFS:
 
 
 remote_params = [S3, GCP, Azure, OSS, SSH, HDFS]
+all_remote_params = [Local] + remote_params
 
 
 @pytest.fixture
@@ -63,15 +70,11 @@ def remote_url(request):
     return request.param.get_url()
 
 
-def pytest_generate_tests(metafunc):
-    if "remote_url" in metafunc.fixturenames:
-        metafunc.parametrize("remote_url", remote_params, indirect=True)
-
-
 def run_dvc(*argv):
     assert main(argv) == 0
 
 
+@pytest.mark.parametrize("remote_url", remote_params, indirect=True)
 def test_get_url(repo_dir, dvc_repo, remote_url):
     run_dvc("remote", "add", "-d", "upstream", remote_url)
     dvc_repo.add(repo_dir.FOO)
@@ -80,6 +83,7 @@ def test_get_url(repo_dir, dvc_repo, remote_url):
     assert api.get_url(repo_dir.FOO) == expected_url
 
 
+@pytest.mark.parametrize("remote_url", remote_params, indirect=True)
 def test_get_url_external(repo_dir, dvc_repo, erepo, remote_url):
     _set_remote_url_and_commit(erepo.dvc, remote_url)
 
@@ -89,6 +93,7 @@ def test_get_url_external(repo_dir, dvc_repo, erepo, remote_url):
     assert api.get_url(repo_dir.FOO, repo=repo_url) == expected_url
 
 
+@pytest.mark.parametrize("remote_url", all_remote_params, indirect=True)
 def test_open(repo_dir, dvc_repo, remote_url):
     run_dvc("remote", "add", "-d", "upstream", remote_url)
     dvc_repo.add(repo_dir.FOO)
@@ -101,6 +106,7 @@ def test_open(repo_dir, dvc_repo, remote_url):
         assert fd.read() == repo_dir.FOO_CONTENTS
 
 
+@pytest.mark.parametrize("remote_url", all_remote_params, indirect=True)
 def test_open_external(repo_dir, dvc_repo, erepo, remote_url):
     erepo.dvc.scm.checkout("branch")
     _set_remote_url_and_commit(erepo.dvc, remote_url)
@@ -120,13 +126,16 @@ def test_open_external(repo_dir, dvc_repo, erepo, remote_url):
     assert api.read("version", repo=repo_url, rev="branch") == "branch"
 
 
-def test_open_missing(erepo):
+@pytest.mark.parametrize("remote_url", all_remote_params, indirect=True)
+def test_missing(repo_dir, dvc_repo, remote_url):
+    run_dvc("add", repo_dir.FOO)
+    run_dvc("remote", "add", "-d", "upstream", remote_url)
+
     # Remove cache to make foo missing
-    shutil.rmtree(erepo.dvc.cache.local.cache_dir)
+    shutil.rmtree(dvc_repo.cache.local.cache_dir)
 
-    repo_url = "file://" + erepo.dvc.root_dir
     with pytest.raises(OutputFileMissingError):
-        api.read(erepo.FOO, repo=repo_url)
+        api.read(repo_dir.FOO)
 
 
 def _set_remote_url_and_commit(repo, remote_url):