diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py index 346aaa925..015297dce 100644 --- a/dvc/dependency/repo.py +++ b/dvc/dependency/repo.py @@ -50,7 +50,7 @@ class RepoDependency(LocalDependency): return external_repo(d["url"], rev=rev) def _get_checksum(self, locked=True): - from dvc.repo.tree import RepoTree + from dvc.tree.repo import RepoTree with self._make_repo(locked=locked) as repo: try: diff --git a/dvc/external_repo.py b/dvc/external_repo.py index c2de3d8e8..965f0acd5 100644 --- a/dvc/external_repo.py +++ b/dvc/external_repo.py @@ -18,10 +18,10 @@ from dvc.exceptions import ( ) from dvc.path_info import PathInfo from dvc.repo import Repo -from dvc.repo.tree import RepoTree from dvc.scm.base import CloneError from dvc.scm.git import Git from dvc.tree.local import LocalTree +from dvc.tree.repo import RepoTree from dvc.utils.fs import remove logger = logging.getLogger(__name__) diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index c518b726f..edebd74fb 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -15,7 +15,7 @@ from dvc.exceptions import ( OutputNotFoundError, ) from dvc.path_info import PathInfo -from dvc.repo.tree import RepoTree +from dvc.tree.repo import RepoTree from dvc.utils.fs import path_isin from ..stage.exceptions import StageFileDoesNotExistError, StageNotFound diff --git a/dvc/repo/experiments/executor.py b/dvc/repo/experiments/executor.py index 3788c3cf4..5f2fd7f46 100644 --- a/dvc/repo/experiments/executor.py +++ b/dvc/repo/experiments/executor.py @@ -7,10 +7,10 @@ from typing import Iterable from funcy import cached_property from dvc.path_info import PathInfo -from dvc.repo.tree import RepoTree from dvc.stage import PipelineStage from dvc.tree.base import BaseTree from dvc.tree.local import LocalTree +from dvc.tree.repo import RepoTree from dvc.utils import relpath from dvc.utils.fs import copy_fobj_to_file, makedirs diff --git a/dvc/repo/ls.py b/dvc/repo/ls.py index 96597fcff..7999bdd8c 100644 --- a/dvc/repo/ls.py +++ b/dvc/repo/ls.py @@ -48,7 +48,7 @@ def ls( def _ls(repo, path_info, recursive=None, dvc_only=False): - from dvc.repo.tree import RepoTree + from dvc.tree.repo import RepoTree def onerror(exc): raise exc diff --git a/dvc/repo/metrics/show.py b/dvc/repo/metrics/show.py index b1b0e9b96..192174073 100644 --- a/dvc/repo/metrics/show.py +++ b/dvc/repo/metrics/show.py @@ -4,7 +4,7 @@ import os from dvc.exceptions import NoMetricsError from dvc.path_info import PathInfo from dvc.repo import locked -from dvc.repo.tree import RepoTree +from dvc.tree.repo import RepoTree from dvc.utils.serialize import YAMLFileCorruptedError, parse_yaml logger = logging.getLogger(__name__) diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py index fd63ec228..a3222b395 100644 --- a/dvc/repo/plots/__init__.py +++ b/dvc/repo/plots/__init__.py @@ -3,8 +3,8 @@ import logging from funcy import first, project from dvc.exceptions import DvcException, NoPlotsError, OutputNotFoundError -from dvc.repo.tree import RepoTree from dvc.schema import PLOT_PROPS +from dvc.tree.repo import RepoTree from dvc.utils import relpath logger = logging.getLogger(__name__) diff --git a/dvc/tree/dvc.py b/dvc/tree/dvc.py new file mode 100644 index 000000000..a88327f8f --- /dev/null +++ b/dvc/tree/dvc.py @@ -0,0 +1,233 @@ +import logging +import os + +from dvc.exceptions import OutputNotFoundError +from dvc.path_info import PathInfo + +from .base import BaseTree, RemoteActionNotImplemented + +logger = logging.getLogger(__name__) + + +class DvcTree(BaseTree): # pylint:disable=abstract-method + """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): + super().__init__(repo, {"url": repo.root_dir}) + self.fetch = fetch + self.stream = stream + + def _find_outs(self, path, *args, **kwargs): + outs = self.repo.find_outs_by_path(path, *args, **kwargs) + + def _is_cached(out): + return out.use_cache + + outs = list(filter(_is_cached, outs)) + if not outs: + raise OutputNotFoundError(path, self.repo) + + return outs + + def _get_granular_checksum(self, path, out, remote=None): + assert isinstance(path, PathInfo) + if not self.fetch and not self.stream: + raise FileNotFoundError + dir_cache = out.get_dir_cache(remote=remote) + for entry in dir_cache: + entry_relpath = entry[out.tree.PARAM_RELPATH] + if os.name == "nt": + entry_relpath = entry_relpath.replace("/", os.sep) + if path == out.path_info / entry_relpath: + return entry[out.tree.PARAM_CHECKSUM] + raise FileNotFoundError + + def open( + self, path, mode="r", encoding="utf-8", remote=None + ): # pylint: disable=arguments-differ + try: + outs = self._find_outs(path, strict=False) + except OutputNotFoundError as exc: + raise FileNotFoundError from exc + + # NOTE: this handles both dirty and checkout-ed out at the same time + if self.repo.tree.exists(path): + return self.repo.tree.open(path, mode=mode, encoding=encoding) + + if len(outs) != 1 or ( + outs[0].is_dir_checksum and path == outs[0].path_info + ): + raise IsADirectoryError + + out = outs[0] + if out.changed_cache(filter_info=path): + if not self.fetch and not self.stream: + raise FileNotFoundError + + remote_obj = self.repo.cloud.get_remote(remote) + if self.stream: + if out.is_dir_checksum: + checksum = self._get_granular_checksum(path, out) + else: + checksum = out.checksum + try: + remote_info = remote_obj.tree.hash_to_path_info(checksum) + return remote_obj.tree.open( + remote_info, mode=mode, encoding=encoding + ) + except RemoteActionNotImplemented: + pass + cache_info = out.get_used_cache(filter_info=path, remote=remote) + self.repo.cloud.pull(cache_info, remote=remote) + + if out.is_dir_checksum: + checksum = self._get_granular_checksum(path, out) + cache_path = out.cache.tree.hash_to_path_info(checksum).url + else: + cache_path = out.cache_path + return open(cache_path, mode=mode, encoding=encoding) + + def exists(self, path): # pylint: disable=arguments-differ + try: + self._find_outs(path, strict=False, recursive=True) + return True + except OutputNotFoundError: + return False + + def isdir(self, path): # pylint: disable=arguments-differ + if not self.exists(path): + return False + + path_info = PathInfo(os.path.abspath(path)) + outs = self._find_outs(path, strict=False, recursive=True) + if len(outs) != 1: + return True + + out = outs[0] + if not out.is_dir_checksum: + return out.path_info != path_info + if out.path_info == path_info: + return True + + # for dir checksum, we need to check if this is a file inside the + # directory + try: + self._get_granular_checksum(path_info, out) + return False + except FileNotFoundError: + return True + + def isfile(self, path): # pylint: disable=arguments-differ + if not self.exists(path): + return False + + return not self.isdir(path) + + def _add_dir(self, top, trie, out, download_callback=None, **kwargs): + if not self.fetch and not self.stream: + return + + # pull dir cache if needed + dir_cache = out.get_dir_cache(**kwargs) + + # pull dir contents if needed + if self.fetch and out.changed_cache(filter_info=top): + used_cache = out.get_used_cache(filter_info=top) + downloaded = self.repo.cloud.pull(used_cache, **kwargs) + if download_callback: + download_callback(downloaded) + + for entry in dir_cache: + entry_relpath = entry[out.tree.PARAM_RELPATH] + if os.name == "nt": + entry_relpath = entry_relpath.replace("/", os.sep) + path_info = out.path_info / entry_relpath + trie[path_info.parts] = None + + def _walk(self, root, trie, topdown=True, **kwargs): + dirs = set() + files = [] + + out = trie.get(root.parts) + if out and out.is_dir_checksum: + self._add_dir(root, trie, out, **kwargs) + + root_len = len(root.parts) + for key, out in trie.iteritems(prefix=root.parts): # noqa: B301 + if key == root.parts: + continue + + name = key[root_len] + if len(key) > root_len + 1 or (out and out.is_dir_checksum): + dirs.add(name) + continue + + files.append(name) + + assert topdown + dirs = list(dirs) + yield root.fspath, dirs, files + + for dname in dirs: + yield from self._walk(root / dname, trie) + + def walk(self, top, topdown=True, onerror=None, **kwargs): + from pygtrie import Trie + + assert topdown + + if not self.exists(top): + if onerror is not None: + onerror(FileNotFoundError(top)) + return + + if not self.isdir(top): + if onerror is not None: + onerror(NotADirectoryError(top)) + return + + root = PathInfo(os.path.abspath(top)) + outs = self._find_outs(top, recursive=True, strict=False) + + trie = Trie() + + for out in outs: + trie[out.path_info.parts] = out + + if out.is_dir_checksum and root.isin_or_eq(out.path_info): + self._add_dir(top, trie, out, **kwargs) + + yield from self._walk(root, trie, topdown=topdown, **kwargs) + + def isdvc(self, path, **kwargs): + try: + return len(self._find_outs(path, **kwargs)) == 1 + except OutputNotFoundError: + pass + return False + + def isexec(self, path): # pylint: disable=unused-argument + return False + + def get_file_hash(self, path_info): + outs = self._find_outs(path_info, strict=False) + if len(outs) != 1: + raise OutputNotFoundError + out = outs[0] + if out.is_dir_checksum: + return ( + out.tree.PARAM_CHECKSUM, + self._get_granular_checksum(path_info, out), + ) + return out.tree.PARAM_CHECKSUM, out.checksum diff --git a/dvc/repo/tree.py b/dvc/tree/repo.py similarity index 60% rename from dvc/repo/tree.py rename to dvc/tree/repo.py index 09bcf286b..4b56dbf9f 100644 --- a/dvc/repo/tree.py +++ b/dvc/tree/repo.py @@ -10,243 +10,22 @@ from pygtrie import StringTrie from dvc.dvcfile import is_valid_filename from dvc.exceptions import OutputNotFoundError from dvc.path_info import PathInfo -from dvc.tree.base import BaseTree, RemoteActionNotImplemented from dvc.utils import file_md5 from dvc.utils.fs import copy_fobj_to_file, makedirs +from .base import BaseTree +from .dvc import DvcTree + if TYPE_CHECKING: from dvc.repo import Repo - from dvc.tree.git import GitTree - from dvc.tree.local import LocalTree + + from .git import GitTree + from .local import LocalTree logger = logging.getLogger(__name__) -class DvcTree(BaseTree): # pylint:disable=abstract-method - """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): - super().__init__(repo, {"url": repo.root_dir}) - self.fetch = fetch - self.stream = stream - - def _find_outs(self, path, *args, **kwargs): - outs = self.repo.find_outs_by_path(path, *args, **kwargs) - - def _is_cached(out): - return out.use_cache - - outs = list(filter(_is_cached, outs)) - if not outs: - raise OutputNotFoundError(path, self.repo) - - return outs - - def _get_granular_checksum(self, path, out, remote=None): - assert isinstance(path, PathInfo) - if not self.fetch and not self.stream: - raise FileNotFoundError - dir_cache = out.get_dir_cache(remote=remote) - for entry in dir_cache: - entry_relpath = entry[out.tree.PARAM_RELPATH] - if os.name == "nt": - entry_relpath = entry_relpath.replace("/", os.sep) - if path == out.path_info / entry_relpath: - return entry[out.tree.PARAM_CHECKSUM] - raise FileNotFoundError - - def open( - self, path, mode="r", encoding="utf-8", remote=None - ): # pylint: disable=arguments-differ - try: - outs = self._find_outs(path, strict=False) - except OutputNotFoundError as exc: - raise FileNotFoundError from exc - - # NOTE: this handles both dirty and checkout-ed out at the same time - if self.repo.tree.exists(path): - return self.repo.tree.open(path, mode=mode, encoding=encoding) - - if len(outs) != 1 or ( - outs[0].is_dir_checksum and path == outs[0].path_info - ): - raise IsADirectoryError - - out = outs[0] - if out.changed_cache(filter_info=path): - if not self.fetch and not self.stream: - raise FileNotFoundError - - remote_obj = self.repo.cloud.get_remote(remote) - if self.stream: - if out.is_dir_checksum: - checksum = self._get_granular_checksum(path, out) - else: - checksum = out.checksum - try: - remote_info = remote_obj.tree.hash_to_path_info(checksum) - return remote_obj.tree.open( - remote_info, mode=mode, encoding=encoding - ) - except RemoteActionNotImplemented: - pass - cache_info = out.get_used_cache(filter_info=path, remote=remote) - self.repo.cloud.pull(cache_info, remote=remote) - - if out.is_dir_checksum: - checksum = self._get_granular_checksum(path, out) - cache_path = out.cache.tree.hash_to_path_info(checksum).url - else: - cache_path = out.cache_path - return open(cache_path, mode=mode, encoding=encoding) - - def exists(self, path): # pylint: disable=arguments-differ - try: - self._find_outs(path, strict=False, recursive=True) - return True - except OutputNotFoundError: - return False - - def isdir(self, path): # pylint: disable=arguments-differ - if not self.exists(path): - return False - - path_info = PathInfo(os.path.abspath(path)) - outs = self._find_outs(path, strict=False, recursive=True) - if len(outs) != 1: - return True - - out = outs[0] - if not out.is_dir_checksum: - return out.path_info != path_info - if out.path_info == path_info: - return True - - # for dir checksum, we need to check if this is a file inside the - # directory - try: - self._get_granular_checksum(path_info, out) - return False - except FileNotFoundError: - return True - - def isfile(self, path): # pylint: disable=arguments-differ - if not self.exists(path): - return False - - return not self.isdir(path) - - def _add_dir(self, top, trie, out, download_callback=None, **kwargs): - if not self.fetch and not self.stream: - return - - # pull dir cache if needed - dir_cache = out.get_dir_cache(**kwargs) - - # pull dir contents if needed - if self.fetch and out.changed_cache(filter_info=top): - used_cache = out.get_used_cache(filter_info=top) - downloaded = self.repo.cloud.pull(used_cache, **kwargs) - if download_callback: - download_callback(downloaded) - - for entry in dir_cache: - entry_relpath = entry[out.tree.PARAM_RELPATH] - if os.name == "nt": - entry_relpath = entry_relpath.replace("/", os.sep) - path_info = out.path_info / entry_relpath - trie[path_info.parts] = None - - def _walk(self, root, trie, topdown=True, **kwargs): - dirs = set() - files = [] - - out = trie.get(root.parts) - if out and out.is_dir_checksum: - self._add_dir(root, trie, out, **kwargs) - - root_len = len(root.parts) - for key, out in trie.iteritems(prefix=root.parts): # noqa: B301 - if key == root.parts: - continue - - name = key[root_len] - if len(key) > root_len + 1 or (out and out.is_dir_checksum): - dirs.add(name) - continue - - files.append(name) - - assert topdown - dirs = list(dirs) - yield root.fspath, dirs, files - - for dname in dirs: - yield from self._walk(root / dname, trie) - - def walk(self, top, topdown=True, onerror=None, **kwargs): - from pygtrie import Trie - - assert topdown - - if not self.exists(top): - if onerror is not None: - onerror(FileNotFoundError(top)) - return - - if not self.isdir(top): - if onerror is not None: - onerror(NotADirectoryError(top)) - return - - root = PathInfo(os.path.abspath(top)) - outs = self._find_outs(top, recursive=True, strict=False) - - trie = Trie() - - for out in outs: - trie[out.path_info.parts] = out - - if out.is_dir_checksum and root.isin_or_eq(out.path_info): - self._add_dir(top, trie, out, **kwargs) - - yield from self._walk(root, trie, topdown=topdown, **kwargs) - - def isdvc(self, path, **kwargs): - try: - return len(self._find_outs(path, **kwargs)) == 1 - except OutputNotFoundError: - pass - return False - - def isexec(self, path): # pylint: disable=unused-argument - return False - - def get_file_hash(self, path_info): - outs = self._find_outs(path_info, strict=False) - if len(outs) != 1: - raise OutputNotFoundError - out = outs[0] - if out.is_dir_checksum: - return ( - out.tree.PARAM_CHECKSUM, - self._get_granular_checksum(path_info, out), - ) - return out.tree.PARAM_CHECKSUM, out.checksum - - class RepoTree(BaseTree): # pylint:disable=abstract-method """DVC + git-tracked files tree. diff --git a/tests/func/test_tree.py b/tests/func/test_tree.py index aec286243..abe002859 100644 --- a/tests/func/test_tree.py +++ b/tests/func/test_tree.py @@ -4,10 +4,10 @@ from os.path import join from dvc.path_info import PathInfo from dvc.repo import Repo -from dvc.repo.tree import RepoTree from dvc.scm import SCM from dvc.tree.git import GitTree from dvc.tree.local import LocalTree +from dvc.tree.repo import RepoTree from dvc.utils.fs import remove from tests.basic_env import TestDir, TestGit, TestGitSubmodule diff --git a/tests/unit/tree/__init__.py b/tests/unit/tree/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/repo/test_tree.py b/tests/unit/tree/test_dvc.py similarity index 99% rename from tests/unit/repo/test_tree.py rename to tests/unit/tree/test_dvc.py index 2897766c5..4c3a428d1 100644 --- a/tests/unit/repo/test_tree.py +++ b/tests/unit/tree/test_dvc.py @@ -4,7 +4,7 @@ import shutil import pytest from dvc.path_info import PathInfo -from dvc.repo.tree import DvcTree +from dvc.tree.dvc import DvcTree def test_exists(tmp_dir, dvc): diff --git a/tests/unit/repo/test_repo_tree.py b/tests/unit/tree/test_repo.py similarity index 99% rename from tests/unit/repo/test_repo_tree.py rename to tests/unit/tree/test_repo.py index aa54a3bd3..087150b9c 100644 --- a/tests/unit/repo/test_repo_tree.py +++ b/tests/unit/tree/test_repo.py @@ -5,7 +5,7 @@ from unittest import mock import pytest from dvc.path_info import PathInfo -from dvc.repo.tree import RepoTree +from dvc.tree.repo import RepoTree def test_exists(tmp_dir, dvc):