mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
repo: Support streaming and pulling files on RepoTree/DvcTree.open() (#3810)
* DvcTree: support stream or fetch on open() * tests: test DvcTree dir cache support for walk() * RepoTree: finish implementation of BaseTree functions * repo: use RepoTree/DvcTree.open() for repo.open_by_relpath * repo: only use RepoTree when needed for now * fix docstring * RepoTree: only use DvcTree.open for DVC outs
This commit is contained in:
Peter Rowlands (변기호)
+2
-2
@@ -339,7 +339,7 @@ class BaseOutput:
|
||||
if self.exists:
|
||||
self.remote.unprotect(self.path_info)
|
||||
|
||||
def _collect_used_dir_cache(
|
||||
def collect_used_dir_cache(
|
||||
self, remote=None, force=False, jobs=None, filter_info=None
|
||||
):
|
||||
"""Get a list of `info`s related to the given directory.
|
||||
@@ -449,7 +449,7 @@ class BaseOutput:
|
||||
return ret
|
||||
|
||||
ret.add_child_cache(
|
||||
self.checksum, self._collect_used_dir_cache(**kwargs),
|
||||
self.checksum, self.collect_used_dir_cache(**kwargs),
|
||||
)
|
||||
|
||||
return ret
|
||||
|
||||
+18
-48
@@ -5,15 +5,12 @@ from functools import wraps
|
||||
from funcy import cached_property, cat, first
|
||||
|
||||
from dvc.config import Config
|
||||
from dvc.exceptions import (
|
||||
FileMissingError,
|
||||
IsADirectoryError,
|
||||
NotDvcRepoError,
|
||||
OutputNotFoundError,
|
||||
)
|
||||
from dvc.exceptions import FileMissingError
|
||||
from dvc.exceptions import IsADirectoryError as DvcIsADirectoryError
|
||||
from dvc.exceptions import NotDvcRepoError, OutputNotFoundError
|
||||
from dvc.ignore import CleanTree
|
||||
from dvc.path_info import PathInfo
|
||||
from dvc.remote.base import RemoteActionNotImplemented
|
||||
from dvc.repo.tree import RepoTree
|
||||
from dvc.utils.fs import path_isin
|
||||
|
||||
from ..utils import parse_target
|
||||
@@ -486,48 +483,21 @@ class Repo:
|
||||
@contextmanager
|
||||
def open_by_relpath(self, path, remote=None, mode="r", encoding=None):
|
||||
"""Opens a specified resource as a file descriptor"""
|
||||
cause = None
|
||||
|
||||
tree = RepoTree(self, stream=True)
|
||||
path = os.path.join(self.root_dir, path)
|
||||
try:
|
||||
out = self.find_out_by_relpath(path)
|
||||
except OutputNotFoundError as exc:
|
||||
out = None
|
||||
cause = exc
|
||||
|
||||
if out and out.use_cache:
|
||||
try:
|
||||
with self._open_cached(out, remote, mode, encoding) as fd:
|
||||
yield fd
|
||||
return
|
||||
except FileNotFoundError as exc:
|
||||
raise FileMissingError(path) from exc
|
||||
|
||||
abs_path = os.path.join(self.root_dir, path)
|
||||
if os.path.exists(abs_path):
|
||||
with open(abs_path, mode=mode, encoding=encoding) as fd:
|
||||
yield fd
|
||||
return
|
||||
|
||||
raise FileMissingError(path) from cause
|
||||
|
||||
def _open_cached(self, out, remote=None, mode="r", encoding=None):
|
||||
if out.isdir():
|
||||
raise IsADirectoryError("Can't open a dir")
|
||||
|
||||
cache_file = self.cache.local.checksum_to_path_info(out.checksum)
|
||||
|
||||
if os.path.exists(cache_file):
|
||||
return open(cache_file, mode=mode, encoding=encoding)
|
||||
|
||||
try:
|
||||
remote_obj = self.cloud.get_remote(remote)
|
||||
remote_info = remote_obj.checksum_to_path_info(out.checksum)
|
||||
return remote_obj.open(remote_info, mode=mode, encoding=encoding)
|
||||
except RemoteActionNotImplemented:
|
||||
with self.state:
|
||||
cache_info = out.get_used_cache(remote=remote)
|
||||
self.cloud.pull(cache_info, remote=remote)
|
||||
|
||||
return open(cache_file, mode=mode, encoding=encoding)
|
||||
with tree.open(
|
||||
os.path.join(self.root_dir, path),
|
||||
mode=mode,
|
||||
encoding=encoding,
|
||||
remote=remote,
|
||||
) as fobj:
|
||||
yield fobj
|
||||
except FileNotFoundError as exc:
|
||||
raise FileMissingError(path) from exc
|
||||
except IsADirectoryError as exc:
|
||||
raise DvcIsADirectoryError from exc
|
||||
|
||||
def close(self):
|
||||
self.scm.close()
|
||||
|
||||
+165
-17
@@ -1,14 +1,32 @@
|
||||
import errno
|
||||
import logging
|
||||
import os
|
||||
|
||||
from dvc.exceptions import OutputNotFoundError
|
||||
from dvc.path_info import PathInfo
|
||||
from dvc.remote.base import RemoteActionNotImplemented
|
||||
from dvc.scm.tree import BaseTree, WorkingTree
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DvcTree(BaseTree):
|
||||
def __init__(self, repo):
|
||||
"""DVC repo tree.
|
||||
|
||||
Args:
|
||||
repo: DVC repo.
|
||||
fetch: if True, uncached DVC outs will be fetched on `open()`.
|
||||
stream: if True, uncached DVC outs will be streamed directly from
|
||||
remote on `open()`.
|
||||
|
||||
`stream` takes precedence over `fetch`. If `stream` is enabled and
|
||||
a remote does not support streaming, uncached DVC outs will be fetched
|
||||
as a fallback.
|
||||
"""
|
||||
|
||||
def __init__(self, repo, fetch=False, stream=False):
|
||||
self.repo = repo
|
||||
self.fetch = fetch
|
||||
self.stream = stream
|
||||
|
||||
def _find_outs(self, path, *args, **kwargs):
|
||||
outs = self.repo.find_outs_by_path(path, *args, **kwargs)
|
||||
@@ -22,14 +40,14 @@ class DvcTree(BaseTree):
|
||||
|
||||
return outs
|
||||
|
||||
def open(self, path, mode="r", encoding="utf-8"):
|
||||
def open(self, path, mode="r", encoding="utf-8", remote=None):
|
||||
try:
|
||||
outs = self._find_outs(path, strict=False)
|
||||
except OutputNotFoundError as exc:
|
||||
raise FileNotFoundError from exc
|
||||
|
||||
if len(outs) != 1 or outs[0].is_dir_checksum:
|
||||
raise OSError(errno.EISDIR)
|
||||
raise IsADirectoryError
|
||||
|
||||
out = outs[0]
|
||||
# temporary hack to make cache use WorkingTree and not GitTree, because
|
||||
@@ -38,7 +56,23 @@ class DvcTree(BaseTree):
|
||||
self.repo.tree = WorkingTree(self.repo.root_dir)
|
||||
try:
|
||||
if out.changed_cache():
|
||||
raise FileNotFoundError
|
||||
if not self.fetch and not self.stream:
|
||||
raise FileNotFoundError
|
||||
|
||||
remote_obj = self.repo.cloud.get_remote(remote)
|
||||
if self.stream:
|
||||
try:
|
||||
remote_info = remote_obj.checksum_to_path_info(
|
||||
out.checksum
|
||||
)
|
||||
return remote_obj.open(
|
||||
remote_info, mode=mode, encoding=encoding
|
||||
)
|
||||
except RemoteActionNotImplemented:
|
||||
pass
|
||||
with self.repo.state:
|
||||
cache_info = out.get_used_cache(remote=remote)
|
||||
self.repo.cloud.pull(cache_info, remote=remote)
|
||||
finally:
|
||||
self.repo.tree = saved_tree
|
||||
|
||||
@@ -78,14 +112,15 @@ class DvcTree(BaseTree):
|
||||
continue
|
||||
|
||||
name = key[root_len]
|
||||
if len(key) > root_len + 1 or out.is_dir_checksum:
|
||||
if len(key) > root_len + 1 or (out and out.is_dir_checksum):
|
||||
dirs.add(name)
|
||||
continue
|
||||
|
||||
files.append(name)
|
||||
|
||||
if topdown:
|
||||
yield root.fspath, list(dirs), files
|
||||
dirs = list(dirs)
|
||||
yield root.fspath, dirs, files
|
||||
|
||||
for dname in dirs:
|
||||
yield from self._walk(root / dname, trie)
|
||||
@@ -111,6 +146,15 @@ class DvcTree(BaseTree):
|
||||
for out in outs:
|
||||
trie[out.path_info.parts] = out
|
||||
|
||||
if out.is_dir_checksum and (self.fetch or self.stream):
|
||||
# will pull dir cache if needed
|
||||
with self.repo.state:
|
||||
cache = out.collect_used_dir_cache()
|
||||
for _, names in cache.scheme_names(out.scheme):
|
||||
for name in names:
|
||||
path_info = out.path_info.parent / name
|
||||
trie[path_info.parts] = None
|
||||
|
||||
yield from self._walk(root, trie, topdown=topdown)
|
||||
|
||||
def isdvc(self, path):
|
||||
@@ -125,17 +169,121 @@ class DvcTree(BaseTree):
|
||||
|
||||
|
||||
class RepoTree(BaseTree):
|
||||
def __init__(self, repo):
|
||||
"""DVC + git-tracked files tree.
|
||||
|
||||
Args:
|
||||
repo: DVC or git repo.
|
||||
|
||||
Any kwargs will be passed to `DvcTree()`.
|
||||
"""
|
||||
|
||||
def __init__(self, repo, **kwargs):
|
||||
self.repo = repo
|
||||
self.dvctree = DvcTree(repo)
|
||||
if hasattr(repo, "dvc_dir"):
|
||||
self.dvctree = DvcTree(repo, **kwargs)
|
||||
else:
|
||||
# git-only erepo's do not need dvctree
|
||||
self.dvctree = None
|
||||
|
||||
def open(self, *args, **kwargs):
|
||||
try:
|
||||
return self.dvctree.open(*args, **kwargs)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
return self.repo.tree.open(*args, **kwargs)
|
||||
def open(self, path, mode="r", encoding="utf-8", **kwargs):
|
||||
if self.dvctree and self.dvctree.exists(path):
|
||||
try:
|
||||
return self.dvctree.open(
|
||||
path, mode=mode, encoding=encoding, **kwargs
|
||||
)
|
||||
except FileNotFoundError:
|
||||
if self.isdvc(path):
|
||||
raise
|
||||
return self.repo.tree.open(path, mode=mode, encoding=encoding)
|
||||
|
||||
def exists(self, path):
|
||||
return self.repo.tree.exists(path) or self.dvctree.exists(path)
|
||||
return self.repo.tree.exists(path) or (
|
||||
self.dvctree and self.dvctree.exists(path)
|
||||
)
|
||||
|
||||
def isdir(self, path):
|
||||
return self.repo.tree.isdir(path) or (
|
||||
self.dvctree and self.dvctree.isdir(path)
|
||||
)
|
||||
|
||||
def isdvc(self, path):
|
||||
return self.dvctree is not None and self.dvctree.isdvc(path)
|
||||
|
||||
def isfile(self, path):
|
||||
return self.repo.tree.isfile(path) or (
|
||||
self.dvctree and self.dvctree.isfile(path)
|
||||
)
|
||||
|
||||
def isexec(self, path):
|
||||
if self.dvctree and self.dvctree.exists(path):
|
||||
return self.dvctree.isexec(path)
|
||||
return self.repo.tree.isexec(path)
|
||||
|
||||
def _walk_one(self, walk):
|
||||
try:
|
||||
root, dirs, files = next(walk)
|
||||
except StopIteration:
|
||||
return
|
||||
yield root, dirs, files
|
||||
for _ in dirs:
|
||||
yield from self._walk_one(walk)
|
||||
|
||||
def _walk(self, dvc_walk, repo_walk):
|
||||
try:
|
||||
_, dvc_dirs, dvc_fnames = next(dvc_walk)
|
||||
repo_root, repo_dirs, repo_fnames = next(repo_walk)
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
# separate subdirs into shared dirs, dvc-only dirs, repo-only dirs
|
||||
dvc_set = set(dvc_dirs)
|
||||
repo_set = set(repo_dirs)
|
||||
dvc_only = list(dvc_set - repo_set)
|
||||
repo_only = list(repo_set - dvc_set)
|
||||
shared = list(dvc_set & repo_set)
|
||||
dirs = shared + dvc_only + repo_only
|
||||
|
||||
# merge file lists
|
||||
files = set(dvc_fnames)
|
||||
for filename in repo_fnames:
|
||||
files.add(filename)
|
||||
|
||||
yield repo_root, dirs, list(files)
|
||||
|
||||
# set dir order for next recursion level - shared dirs first so that
|
||||
# next() for both generators recurses into the same shared directory
|
||||
dvc_dirs[:] = [dirname for dirname in dirs if dirname in dvc_set]
|
||||
repo_dirs[:] = [dirname for dirname in dirs if dirname in repo_set]
|
||||
|
||||
for dirname in dirs:
|
||||
if dirname in shared:
|
||||
yield from self._walk(dvc_walk, repo_walk)
|
||||
elif dirname in dvc_set:
|
||||
yield from self._walk_one(dvc_walk)
|
||||
elif dirname in repo_set:
|
||||
yield from self._walk_one(repo_walk)
|
||||
|
||||
def walk(self, top, topdown=True):
|
||||
"""Walk and merge both DVC and repo trees."""
|
||||
assert topdown
|
||||
|
||||
if not self.exists(top):
|
||||
raise FileNotFoundError
|
||||
|
||||
if not self.isdir(top):
|
||||
raise NotADirectoryError
|
||||
|
||||
dvc_exists = self.dvctree and self.dvctree.exists(top)
|
||||
repo_exists = self.repo.tree.exists(top)
|
||||
if dvc_exists and not repo_exists:
|
||||
yield from self.dvctree.walk(top, topdown=topdown)
|
||||
return
|
||||
if repo_exists and not dvc_exists:
|
||||
yield from self.repo.tree.walk(top, topdown=topdown)
|
||||
return
|
||||
if not dvc_exists and not repo_exists:
|
||||
raise FileNotFoundError
|
||||
|
||||
dvc_walk = self.dvctree.walk(top, topdown=topdown)
|
||||
repo_walk = self.repo.tree.walk(top, topdown=topdown)
|
||||
yield from self._walk(dvc_walk, repo_walk)
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
import os
|
||||
import shutil
|
||||
|
||||
from dvc.repo.tree import RepoTree
|
||||
|
||||
|
||||
def test_exists(tmp_dir, dvc):
|
||||
tmp_dir.gen("foo", "foo")
|
||||
dvc.add("foo")
|
||||
(tmp_dir / "foo").unlink()
|
||||
|
||||
tree = RepoTree(dvc)
|
||||
assert tree.exists("foo")
|
||||
|
||||
|
||||
def test_open(tmp_dir, dvc):
|
||||
tmp_dir.gen("foo", "foo")
|
||||
dvc.add("foo")
|
||||
(tmp_dir / "foo").unlink()
|
||||
|
||||
tree = RepoTree(dvc)
|
||||
with tree.open("foo", "r") as fobj:
|
||||
assert fobj.read() == "foo"
|
||||
|
||||
|
||||
def test_open_in_history(tmp_dir, scm, dvc):
|
||||
tmp_dir.gen("foo", "foo")
|
||||
dvc.add("foo")
|
||||
dvc.scm.add(["foo.dvc", ".gitignore"])
|
||||
dvc.scm.commit("foo")
|
||||
|
||||
tmp_dir.gen("foo", "foofoo")
|
||||
dvc.add("foo")
|
||||
dvc.scm.add(["foo.dvc", ".gitignore"])
|
||||
dvc.scm.commit("foofoo")
|
||||
|
||||
for rev in dvc.brancher(revs=["HEAD~1"]):
|
||||
if rev == "working tree":
|
||||
continue
|
||||
|
||||
tree = RepoTree(dvc)
|
||||
with tree.open("foo", "r") as fobj:
|
||||
assert fobj.read() == "foo"
|
||||
|
||||
|
||||
def test_isdir_isfile(tmp_dir, dvc):
|
||||
tmp_dir.gen({"datafile": "data", "datadir": {"foo": "foo", "bar": "bar"}})
|
||||
|
||||
tree = RepoTree(dvc)
|
||||
assert tree.isdir("datadir")
|
||||
assert not tree.isfile("datadir")
|
||||
assert not tree.isdvc("datadir")
|
||||
assert not tree.isdir("datafile")
|
||||
assert tree.isfile("datafile")
|
||||
assert not tree.isdvc("datafile")
|
||||
|
||||
dvc.add(["datadir", "datafile"])
|
||||
shutil.rmtree(tmp_dir / "datadir")
|
||||
(tmp_dir / "datafile").unlink()
|
||||
|
||||
assert tree.isdir("datadir")
|
||||
assert not tree.isfile("datadir")
|
||||
assert tree.isdvc("datadir")
|
||||
assert not tree.isdir("datafile")
|
||||
assert tree.isfile("datafile")
|
||||
assert tree.isdvc("datafile")
|
||||
|
||||
|
||||
def test_isdir_mixed(tmp_dir, dvc):
|
||||
tmp_dir.gen({"dir": {"foo": "foo", "bar": "bar"}})
|
||||
|
||||
dvc.add(str(tmp_dir / "dir" / "foo"))
|
||||
|
||||
tree = RepoTree(dvc)
|
||||
assert tree.isdir("dir")
|
||||
assert not tree.isfile("dir")
|
||||
|
||||
|
||||
def test_walk(tmp_dir, dvc):
|
||||
tmp_dir.gen(
|
||||
{
|
||||
"dir": {
|
||||
"subdir1": {"foo1": "foo1", "bar1": "bar1"},
|
||||
"subdir2": {"foo2": "foo2"},
|
||||
}
|
||||
}
|
||||
)
|
||||
dvc.add(str(tmp_dir / "dir"), recursive=True)
|
||||
tmp_dir.gen({"dir": {"foo": "foo", "bar": "bar"}})
|
||||
tree = RepoTree(dvc)
|
||||
|
||||
expected = [
|
||||
os.path.join("dir", "subdir1"),
|
||||
os.path.join("dir", "subdir2"),
|
||||
os.path.join("dir", "subdir1", "foo1"),
|
||||
os.path.join("dir", "subdir1", "foo1.dvc"),
|
||||
os.path.join("dir", "subdir1", "bar1"),
|
||||
os.path.join("dir", "subdir1", "bar1.dvc"),
|
||||
os.path.join("dir", "subdir2", "foo2"),
|
||||
os.path.join("dir", "subdir2", "foo2.dvc"),
|
||||
os.path.join("dir", "foo"),
|
||||
os.path.join("dir", "bar"),
|
||||
]
|
||||
|
||||
actual = []
|
||||
for root, dirs, files in tree.walk("dir"):
|
||||
for entry in dirs + files:
|
||||
actual.append(os.path.join(root, entry))
|
||||
|
||||
assert set(actual) == set(expected)
|
||||
assert len(actual) == len(expected)
|
||||
|
||||
|
||||
def test_isdvc(tmp_dir, dvc):
|
||||
tmp_dir.gen({"foo": "foo", "bar": "bar"})
|
||||
dvc.add("foo")
|
||||
tree = RepoTree(dvc)
|
||||
assert tree.isdvc("foo")
|
||||
assert not tree.isdvc("bar")
|
||||
@@ -1,6 +1,8 @@
|
||||
import os
|
||||
import shutil
|
||||
|
||||
import pytest
|
||||
|
||||
from dvc.repo.tree import DvcTree
|
||||
|
||||
|
||||
@@ -106,6 +108,50 @@ def test_walk(tmp_dir, dvc):
|
||||
assert len(actual) == len(expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fetch,expected",
|
||||
[
|
||||
(False, []),
|
||||
(
|
||||
True,
|
||||
[
|
||||
"dir/subdir1",
|
||||
"dir/subdir2",
|
||||
"dir/subdir1/foo1",
|
||||
"dir/subdir1/bar1",
|
||||
"dir/subdir2/foo2",
|
||||
"dir/foo",
|
||||
"dir/bar",
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_walk_dir(tmp_dir, dvc, fetch, expected):
|
||||
tmp_dir.gen(
|
||||
{
|
||||
"dir": {
|
||||
"subdir1": {"foo1": "foo1", "bar1": "bar1"},
|
||||
"subdir2": {"foo2": "foo2"},
|
||||
"foo": "foo",
|
||||
"bar": "bar",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
dvc.add("dir")
|
||||
tree = DvcTree(dvc, fetch=fetch)
|
||||
|
||||
expected = [str(tmp_dir / path) for path in expected]
|
||||
|
||||
actual = []
|
||||
for root, dirs, files in tree.walk("dir"):
|
||||
for entry in dirs + files:
|
||||
actual.append(os.path.join(root, entry))
|
||||
|
||||
assert set(actual) == set(expected)
|
||||
assert len(actual) == len(expected)
|
||||
|
||||
|
||||
def test_isdvc(tmp_dir, dvc):
|
||||
tmp_dir.gen({"foo": "foo", "bar": "bar"})
|
||||
dvc.add("foo")
|
||||
|
||||
Reference in New Issue
Block a user