From 3b23dfebaed5fc245dfaf15a3f5952d485d3ebf9 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Wed, 9 Dec 2020 23:51:23 +0200 Subject: [PATCH] git: use backend-provided trie in GitTree (#5044) --- dvc/scm/git/__init__.py | 11 +- dvc/scm/git/backend/base.py | 5 +- dvc/scm/git/backend/dulwich.py | 48 +++++- dvc/scm/git/backend/gitpython.py | 43 +++++- dvc/scm/git/objects.py | 120 +++++++++++++++ dvc/tree/git.py | 248 ++++++++++--------------------- tests/func/test_tree.py | 13 +- 7 files changed, 296 insertions(+), 192 deletions(-) create mode 100644 dvc/scm/git/objects.py diff --git a/dvc/scm/git/__init__.py b/dvc/scm/git/__init__.py index 27472e3bf..40107f139 100644 --- a/dvc/scm/git/__init__.py +++ b/dvc/scm/git/__init__.py @@ -296,6 +296,16 @@ class Git(Base): pass raise NoGitBackendError(name) + def get_tree(self, rev: str, **kwargs): + from dvc.tree.git import GitTree + + from .objects import GitTrie + + resolved = self.resolve_rev(rev) + tree_obj = self._backend_func("get_tree_obj", rev=resolved) + trie = GitTrie(tree_obj, resolved) + return GitTree(self.root_dir, trie, **kwargs) + is_ignored = partialmethod(_backend_func, "is_ignored") add = partialmethod(_backend_func, "add") commit = partialmethod(_backend_func, "commit") @@ -311,7 +321,6 @@ class Git(Base): list_branches = partialmethod(_backend_func, "list_branches") list_tags = partialmethod(_backend_func, "list_tags") list_all_commits = partialmethod(_backend_func, "list_all_commits") - get_tree = partialmethod(_backend_func, "get_tree") get_rev = partialmethod(_backend_func, "get_rev") _resolve_rev = partialmethod(_backend_func, "resolve_rev") resolve_commit = partialmethod(_backend_func, "resolve_commit") diff --git a/dvc/scm/git/backend/base.py b/dvc/scm/git/backend/base.py index 5b1d8a03c..3f61b41c3 100644 --- a/dvc/scm/git/backend/base.py +++ b/dvc/scm/git/backend/base.py @@ -3,7 +3,8 @@ from abc import ABC, abstractmethod from typing import Callable, Iterable, Optional, Tuple from dvc.scm.base import SCMError -from dvc.tree.base import BaseTree + +from ..objects import GitObject class NoGitBackendError(SCMError): @@ -109,7 +110,7 @@ class BaseGitBackend(ABC): pass @abstractmethod - def get_tree(self, rev: str, **kwargs) -> BaseTree: + def get_tree_obj(self, rev: str, **kwargs) -> GitObject: pass @abstractmethod diff --git a/dvc/scm/git/backend/dulwich.py b/dvc/scm/git/backend/dulwich.py index 478ea23cd..8a310aa4e 100644 --- a/dvc/scm/git/backend/dulwich.py +++ b/dvc/scm/git/backend/dulwich.py @@ -1,18 +1,55 @@ import fnmatch +import locale import logging import os -from io import BytesIO +import stat +from io import BytesIO, StringIO from typing import Callable, Iterable, Optional, Tuple from dvc.scm.base import SCMError -from dvc.tree.base import BaseTree from dvc.utils import relpath +from ..objects import GitObject from .base import BaseGitBackend logger = logging.getLogger(__name__) +class DulwichObject(GitObject): + def __init__(self, repo, name, mode, sha): + self.repo = repo + self._name = name + self._mode = mode + self.sha = sha + + def open(self, mode: str = "r", encoding: str = None): + if not encoding: + encoding = locale.getpreferredencoding(False) + # NOTE: we didn't load the object before as Dulwich will also try to + # load the contents of it into memory, which will slow down Trie + # building considerably. + obj = self.repo[self.sha] + data = obj.as_raw_string() + if mode == "rb": + return BytesIO(data) + return StringIO(data.decode(encoding)) + + @property + def name(self) -> str: + return self._name + + @property + def mode(self) -> int: + return self._mode + + def scandir(self) -> Iterable["DulwichObject"]: + tree = self.repo[self.sha] + for entry in tree.iteritems(): # noqa: B301 + yield DulwichObject( + self.repo, entry.path.decode(), entry.mode, entry.sha + ) + + class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method """Dulwich Git backend.""" @@ -125,8 +162,11 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method def list_all_commits(self) -> Iterable[str]: raise NotImplementedError - def get_tree(self, rev: str, **kwargs) -> BaseTree: - raise NotImplementedError + def get_tree_obj(self, rev: str, **kwargs) -> DulwichObject: + from dulwich.objectspec import parse_tree + + tree = parse_tree(self.repo, rev) + return DulwichObject(self.repo, ".", stat.S_IFDIR, tree.id) def get_rev(self) -> str: raise NotImplementedError diff --git a/dvc/scm/git/backend/gitpython.py b/dvc/scm/git/backend/gitpython.py index 2d613840d..c2002905f 100644 --- a/dvc/scm/git/backend/gitpython.py +++ b/dvc/scm/git/backend/gitpython.py @@ -1,3 +1,5 @@ +import io +import locale import logging import os from functools import partial @@ -7,9 +9,9 @@ from funcy import first from dvc.progress import Tqdm from dvc.scm.base import CloneError, RevError, SCMError -from dvc.tree.base import BaseTree from dvc.utils import fix_env, is_binary +from ..objects import GitObject from .base import BaseGitBackend logger = logging.getLogger(__name__) @@ -47,6 +49,38 @@ class TqdmGit(Tqdm): return ops.get(op_code & OP.OP_MASK, "") +class GitPythonObject(GitObject): + def __init__(self, obj): + self.obj = obj + + def open(self, mode: str = "r", encoding: str = None): + if not encoding: + encoding = locale.getpreferredencoding(False) + # GitPython's obj.data_stream is a fragile thing, it is better to + # read it immediately, also it needs to be to decoded if we follow + # the `open()` behavior (since data_stream.read() returns bytes, + # and `open` with default "r" mode returns str) + data = self.obj.data_stream.read() + if mode == "rb": + return io.BytesIO(data) + return io.StringIO(data.decode(encoding)) + + @property + def name(self) -> str: + # NOTE: `obj.name` is not always a basename. See [1] for more details. + # + # [1] https://github.com/iterative/dvc/issues/3481 + return os.path.basename(self.obj.path) + + @property + def mode(self) -> int: + return self.obj.mode + + def scandir(self) -> Iterable["GitPythonObject"]: + for obj in self.obj: + yield GitPythonObject(obj) + + class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method """git-python Git backend.""" @@ -227,10 +261,9 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method ) ] - def get_tree(self, rev: str, **kwargs) -> BaseTree: - from dvc.tree.git import GitTree - - return GitTree(self.repo, self.resolve_rev(rev), **kwargs) + def get_tree_obj(self, rev: str, **kwargs) -> GitPythonObject: + tree = self.repo.tree(rev) + return GitPythonObject(tree) def get_rev(self): return self.repo.rev_parse("HEAD").hexsha diff --git a/dvc/scm/git/objects.py b/dvc/scm/git/objects.py new file mode 100644 index 000000000..990e5ffd2 --- /dev/null +++ b/dvc/scm/git/objects.py @@ -0,0 +1,120 @@ +import os +import stat +from abc import ABC, abstractmethod +from typing import Iterable, Optional + +from pygtrie import Trie + +S_IFGITLINK = 0o160000 + + +def S_ISGITLINK(m: int) -> bool: + return stat.S_IFMT(m) == S_IFGITLINK + + +class GitObject(ABC): + @abstractmethod + def open(self, mode: str = "r", encoding: str = None): + pass + + @property + @abstractmethod + def name(self) -> str: + pass + + @property + @abstractmethod + def mode(self) -> int: + pass + + @abstractmethod + def scandir(self) -> Iterable["GitObject"]: + pass + + @property + def isfile(self) -> bool: + return stat.S_ISREG(self.mode) + + @property + def isdir(self) -> bool: + return stat.S_ISDIR(self.mode) + + @property + def issubmodule(self) -> bool: + return S_ISGITLINK(self.mode) + + +class GitTrie: + def __init__(self, tree: GitObject, rev: str): + self.tree = tree + self.rev = rev + self.trie = Trie() + + self.trie[()] = tree + self._build(tree, ()) + + def _build(self, tree: GitObject, path: tuple): + for obj in tree.scandir(): + obj_path = path + (obj.name,) + self.trie[obj_path] = obj + + if obj.isdir: + self._build(obj, obj_path) + + def open( + self, + key: tuple, + mode: Optional[str] = "r", + encoding: Optional[str] = None, + ): + obj = self.trie[key] + if obj.isdir: + raise IsADirectoryError + + return obj.open(mode=mode, encoding=encoding) + + def exists(self, key: tuple) -> bool: + return bool(self.trie.has_node(key)) + + def isdir(self, key: tuple) -> bool: + try: + obj = self.trie[key] + except KeyError: + return False + return obj.isdir + + def isfile(self, key: tuple) -> bool: + try: + obj = self.trie[key] + except KeyError: + return False + + return obj.isfile + + def walk(self, top: tuple, topdown: Optional[bool] = True): + dirs = [] + nondirs = [] + + def node_factory(_, path, children, obj): + if path == top: + assert obj.isdir + list(filter(None, children)) + elif obj.isdir: + dirs.append(obj.name) + else: + nondirs.append(obj.name) + + self.trie.traverse(node_factory, prefix=top) + + if topdown: + yield top, dirs, nondirs + + for dname in dirs: + yield from self.walk(top + (dname,), topdown=topdown) + + if not topdown: + yield top, dirs, nondirs + + def stat(self, key: tuple) -> os.stat_result: + obj = self.trie[key] + return os.stat_result((obj.mode, 0, 0, 0, 0, 0, 0, 0, 0, 0)) diff --git a/dvc/tree/git.py b/dvc/tree/git.py index 34faec2d9..6751404f9 100644 --- a/dvc/tree/git.py +++ b/dvc/tree/git.py @@ -1,69 +1,43 @@ import errno -import io import os -from functools import lru_cache from funcy import cached_property -from dvc.exceptions import DvcException from dvc.tree.base import BaseTree from dvc.utils import is_exec, relpath -# see git-fast-import(1) -GIT_MODE_DIR = 0o40000 -GIT_MODE_FILE = 0o644 - - -def _item_basename(item): - # NOTE: `item.name` is not always a basename. See [1] for more details. - # - # [1] https://github.com/iterative/dvc/issues/3481#issuecomment-600693884 - return os.path.basename(item.path) - - -def _is_tree_and_contains(obj, path): - if obj.mode != GIT_MODE_DIR: - return False - # see https://github.com/gitpython-developers/GitPython/issues/851 - # `return (i in tree)` doesn't work so here is a workaround: - for item in obj: - if _item_basename(item) == path: - return True - return False - - -# NOTE: small lru cache covers sequential exists() calls -@lru_cache(maxsize=1) -def _get_obj(tree, path): - if not path or path == ".": - return tree - for i in path.split(os.sep): - if not _is_tree_and_contains(tree, i): - # there is no tree for specified path - return None - tree = tree[i] - return tree - class GitTree(BaseTree): # pylint:disable=abstract-method """Proxies the repo file access methods to Git objects""" - def __init__(self, git, rev, use_dvcignore=False, dvcignore_root=None): - """Create GitTree instance - - Args: - git (dvc.scm.Git): - branch: - """ + def __init__( + self, root_dir, trie, use_dvcignore=False, dvcignore_root=None + ): super().__init__(None, {}) - self.git = git - self.rev = rev + self._tree_root = root_dir + self.trie = trie self.use_dvcignore = use_dvcignore self.dvcignore_root = dvcignore_root + @property + def rev(self): + return self.trie.rev + @property def tree_root(self): - return self.git.working_dir + return self._tree_root + + def _get_key(self, path): + if isinstance(path, str): + if not os.path.isabs(path): + relparts = path.split(os.sep) + else: + relparts = relpath(path, self.tree_root).split(os.sep) + else: + relparts = path.relative_to(self.tree_root).parts + if relparts == ["."]: + return () + return tuple(relparts) @cached_property def dvcignore(self): @@ -76,96 +50,48 @@ class GitTree(BaseTree): # pylint:disable=abstract-method def open( self, path, mode="r", encoding=None ): # pylint: disable=arguments-differ - assert mode in {"r", "rb"} + # NOTE: this is incorrect, we need to use locale to determine default + # encoding. encoding = encoding or "utf-8" - relative_path = relpath(path, self.git.working_dir) - - obj = self._git_object_by_path(path) - if obj is None: - msg = f"No such file in branch '{self.rev}'" - raise OSError(errno.ENOENT, msg, relative_path) - if obj.mode == GIT_MODE_DIR: - raise OSError(errno.EISDIR, "Is a directory", relative_path) - - # GitPython's obj.data_stream is a fragile thing, it is better to - # read it immediately, also it needs to be to decoded if we follow - # the `open()` behavior (since data_stream.read() returns bytes, - # and `open` with default "r" mode returns str) - data = obj.data_stream.read() - if mode == "rb": - return io.BytesIO(data) - return io.StringIO(data.decode(encoding)) + key = self._get_key(path) + try: + return self.trie.open(key, mode=mode, encoding=encoding) + except KeyError as exc: + msg = os.strerror(errno.ENOENT) + f"in branch '{self.rev}'" + raise FileNotFoundError(errno.ENOENT, msg, path) from exc + except IsADirectoryError as exc: + raise IsADirectoryError( + errno.EISDIR, os.strerror(errno.EISDIR), path + ) from exc def exists( self, path, use_dvcignore=True ): # pylint: disable=arguments-differ - if self._git_object_by_path(path) is None: - return False - if not use_dvcignore: - return True + def _is_ignored(path): + return self.dvcignore.is_ignored_file( + path + ) or self.dvcignore.is_ignored_dir(path) - return not self.dvcignore.is_ignored_file( - path - ) and not self.dvcignore.is_ignored_dir(path) + if use_dvcignore and _is_ignored(path): + return False + + key = self._get_key(path) + return self.trie.exists(key) def isdir( self, path, use_dvcignore=True ): # pylint: disable=arguments-differ - obj = self._git_object_by_path(path) - if obj is None: + if use_dvcignore and self.dvcignore.is_ignored_dir(path): return False - if obj.mode != GIT_MODE_DIR: - return False - return not (use_dvcignore and self.dvcignore.is_ignored_dir(path)) + key = self._get_key(path) + return self.trie.isdir(key) def isfile(self, path): # pylint: disable=arguments-differ - obj = self._git_object_by_path(path) - if obj is None: + if self.dvcignore.is_ignored_file(path): return False - # according to git-fast-import(1) file mode could be 644 or 755 - if obj.mode & GIT_MODE_FILE != GIT_MODE_FILE: - return False - return not self.dvcignore.is_ignored_file(path) - - @cached_property - def _tree(self): - import git - - try: - return self.git.tree(self.rev) - except git.exc.BadName as exc: # pylint: disable=no-member - raise DvcException( - "revision '{}' not found in Git '{}'".format( - self.rev, os.path.relpath(self.git.working_dir) - ) - ) from exc - - def _git_object_by_path(self, path): - path = relpath(os.path.realpath(path), self.git.working_dir) - if path.split(os.sep, 1)[0] == "..": - # path points outside of git repository - return None - - return _get_obj(self._tree, path) - - def _walk(self, tree, topdown=True): - dirs, nondirs = [], [] - for item in tree: - name = _item_basename(item) - if item.mode == GIT_MODE_DIR: - dirs.append(name) - else: - nondirs.append(name) - - if topdown: - yield os.path.normpath(tree.abspath), dirs, nondirs - - for i in dirs: - yield from self._walk(tree[i], topdown=topdown) - - if not topdown: - yield os.path.normpath(tree.abspath), dirs, nondirs + key = self._get_key(path) + return self.trie.isfile(key) def walk( self, @@ -180,66 +106,42 @@ class GitTree(BaseTree): # pylint:disable=abstract-method See `os.walk` for the docs. Differences: - no support for symlinks """ + if not self.isdir(top, use_dvcignore=use_dvcignore): + if onerror: + if self.exists(top, use_dvcignore=use_dvcignore): + exc = NotADirectoryError( + errno.ENOTDIR, os.strerror(errno.ENOTDIR), top + ) + else: + exc = FileNotFoundError( + errno.ENOENT, os.strerror(errno.ENOENT), top + ) + onerror(exc) + return [] - tree = self._git_object_by_path(top) - if tree is None: - if onerror is not None: - onerror(OSError(errno.ENOENT, "No such file", top)) - return - if tree.mode != GIT_MODE_DIR: - if onerror is not None: - onerror(NotADirectoryError(top)) - return - - for root, dirs, files in self._walk(tree, topdown=topdown): + key = self._get_key(top) + for prefix, dirs, files in self.trie.walk(key, topdown=topdown): + if prefix: + root = os.path.join(self.tree_root, os.sep.join(prefix)) + else: + root = self.tree_root if use_dvcignore: dirs[:], files[:] = self.dvcignore( - os.path.abspath(root), - dirs, - files, - ignore_subrepos=ignore_subrepos, + root, dirs, files, ignore_subrepos=ignore_subrepos, ) yield root, dirs, files def isexec(self, path): - if not self.exists(path): - return False - - mode = self.stat(path).st_mode - return is_exec(mode) + return is_exec(self.stat(path).st_mode) def stat(self, path): - import git - - def to_ctime(git_time): - sec, nano_sec = git_time - return sec + nano_sec / 1000000000 - - if self.dvcignore.is_ignored(path): - raise FileNotFoundError - - obj = self._git_object_by_path(path) - if obj is None: - raise OSError(errno.ENOENT, "No such file") - entry = git.index.IndexEntry.from_blob(obj) - - # os.stat_result takes a tuple in the form: - # (mode, ino, dev, nlink, uid, gid, size, atime, mtime, ctime) - return os.stat_result( - ( - entry.mode, - entry.inode, - entry.dev, - 0, - entry.uid, - entry.gid, - entry.size, - # git index has no atime equivalent, use mtime - to_ctime(entry.mtime), - to_ctime(entry.mtime), - to_ctime(entry.ctime), + key = self._get_key(path) + try: + return self.trie.stat(key) + except KeyError: + raise FileNotFoundError( + errno.ENOENT, os.strerror(errno.ENOENT), path ) - ) @property def hash_jobs(self): # pylint: disable=invalid-overridden-method diff --git a/tests/func/test_tree.py b/tests/func/test_tree.py index a17adad4c..7bb5d015f 100644 --- a/tests/func/test_tree.py +++ b/tests/func/test_tree.py @@ -5,7 +5,6 @@ from os.path import join from dvc.path_info import PathInfo from dvc.repo import Repo 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 @@ -45,7 +44,7 @@ class GitTreeTests: self.scm.add([self.FOO, self.UNICODE, self.DATA_DIR]) self.scm.commit("add") - tree = GitTree(self.git, "master") + tree = self.scm.get_tree("master") with tree.open(self.FOO) as fd: self.assertEqual(fd.read(), self.FOO_CONTENTS) with tree.open(self.UNICODE) as fd: @@ -56,14 +55,14 @@ class GitTreeTests: tree.open(self.DATA_DIR) def test_exists(self): - tree = GitTree(self.git, "master") + tree = self.scm.get_tree("master") self.assertFalse(tree.exists(self.FOO)) self.assertFalse(tree.exists(self.UNICODE)) self.assertFalse(tree.exists(self.DATA_DIR)) self.scm.add([self.FOO, self.UNICODE, self.DATA]) self.scm.commit("add") - tree = GitTree(self.git, "master") + tree = self.scm.get_tree("master") self.assertTrue(tree.exists(self.FOO)) self.assertTrue(tree.exists(self.UNICODE)) self.assertTrue(tree.exists(self.DATA_DIR)) @@ -73,7 +72,7 @@ class GitTreeTests: self.scm.add([self.FOO, self.DATA_DIR]) self.scm.commit("add") - tree = GitTree(self.git, "master") + tree = self.scm.get_tree("master") self.assertTrue(tree.isdir(self.DATA_DIR)) self.assertFalse(tree.isdir(self.FOO)) self.assertFalse(tree.isdir("non-existing-file")) @@ -82,7 +81,7 @@ class GitTreeTests: self.scm.add([self.FOO, self.DATA_DIR]) self.scm.commit("add") - tree = GitTree(self.git, "master") + tree = self.scm.get_tree("master") self.assertTrue(tree.isfile(self.FOO)) self.assertFalse(tree.isfile(self.DATA_DIR)) self.assertFalse(tree.isfile("not-existing-file")) @@ -162,7 +161,7 @@ class TestWalkInGit(AssertWalkEqualMixin, TestGit): scm = SCM(self._root_dir) scm.add([self.DATA_SUB_DIR]) scm.commit("add data_dir/data_sub_dir/data_sub") - tree = GitTree(self.git, "master") + tree = scm.get_tree("master") self.assertWalkEqual( tree.walk("."), [