metadata() on RepoTree (#4434)

* WIP: metadata api on RepoTree

* Make outs optional

* add tests + misc. changes

* reorder tests and fields

dvc/repo/ls.py

@@ -1,3 +1,5 @@
+from itertools import chain
+
 from dvc.exceptions import PathMissingError
 from dvc.path_info import PathInfo
 
@@ -57,48 +59,32 @@ def _ls(repo, path_info, recursive=None, dvc_only=False):
     # fetch directory listings, but don't want to fetch file contents.
     tree = RepoTree(repo, stream=True)
 
-    ret = {}
+    infos = []
     try:
         for root, dirs, files in tree.walk(
             path_info.fspath, onerror=onerror, dvcfiles=True
         ):
-            for fname in files:
-                info = PathInfo(root) / fname
-                dvc = tree.isdvc(info)
-                if dvc or not dvc_only:
-                    path = str(info.relative_to(path_info))
-                    ret[path] = {
-                        "isout": dvc,
-                        "isdir": False,
-                        "isexec": False if dvc else tree.isexec(info),
-                    }
-
+            entries = chain(files, dirs) if not recursive else files
+            infos.extend(PathInfo(root) / entry for entry in entries)
             if not recursive:
-                for dname in dirs:
-                    info = PathInfo(root) / dname
-                    # pylint:disable=protected-access
-                    _, dvctree = tree._get_tree_pair(info)  # noqa
-                    if not dvc_only or (dvctree and dvctree.exists(info)):
-                        dvc = tree.isdvc(info)
-                        path = str(info.relative_to(path_info))
-                        ret[path] = {
-                            "isout": dvc,
-                            "isdir": True,
-                            "isexec": False if dvc else tree.isexec(info),
-                        }
                 break
     except NotADirectoryError:
-        dvc = tree.isdvc(path_info)
-        if dvc or not dvc_only:
-            return {
-                path_info.name: {
-                    "isout": dvc,
-                    "isdir": False,
-                    "isexec": False if dvc else tree.isexec(path_info),
-                }
-            }
-        return {}
+        infos.append(path_info)
     except FileNotFoundError:
         return {}
 
+    ret = {}
+    for info in infos:
+        metadata = tree.metadata(info)
+        if metadata.output_exists or not dvc_only:
+            path = (
+                path_info.name
+                if path_info == info
+                else str(info.relative_to(path_info))
+            )
+            ret[path] = {
+                "isout": metadata.is_output,
+                "isdir": metadata.isdir,
+                "isexec": metadata.is_exec,
+            }
     return ret

dvc/tree/_metadata.py

@@ -0,0 +1,68 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from dvc.output import BaseOutput
+from dvc.path_info import PathInfo
+
+
+@dataclass
+class Metadata:
+    """
+    Container for storing metadata for a given path, similar in spirit to
+    `os.stat_result`.
+    """
+
+    # required field
+    path_info: PathInfo
+
+    # computed fields
+    is_output: bool = field(init=False, default=False)  # is it an output?
+    # is the path part of an output?
+    part_of_output: bool = field(init=False, default=False)
+    # does the path contain outputs?
+    contains_outputs: bool = field(init=False, default=False)
+    # is the path tracked by dvc? equivalent to (is_output or part_of_output)
+    is_dvc: bool = field(init=False, default=False)
+    # does dvc has any output in that path or inside it? equiv. to bool(outs)
+    output_exists: bool = field(init=False)
+
+    # optional fields
+    isdir: bool = False  # is it a directory?
+    is_exec: bool = False  # is it an executable?
+    outs: List[BaseOutput] = field(default_factory=list)  # list of outputs
+
+    def __post_init__(self):
+        self.output_exists = bool(self.outs)
+
+        if not self.output_exists:
+            return
+
+        # it can contain multiple outputs
+        if len(self.outs) > 1:
+            self.contains_outputs = True
+        else:
+            out = self.outs[0]
+            # or, the path itself could be an output
+            self.is_output = self.path_info == out.path_info
+
+            # or, a directory must have an output somewhere deep inside it
+            if not self.is_output:
+                self.contains_outputs = out.path_info.isin(self.path_info)
+
+        # or, the path could be a part of an output, i.e. inside of an output
+        self.part_of_output = not self.is_output and not self.contains_outputs
+
+        self.is_dvc = self.is_output or self.part_of_output
+        # if it contains outputs, it must be a dir
+        self.isdir = self.contains_outputs
+
+    @property
+    def isfile(self):
+        # self.isdir might require additional computation, especially if it's
+        # a dir that's inside a dvc-tracked dir
+        # since, `Metadata()` does not support being used as null (i.e. when
+        # path does not exist), `isfile` becomes `isdir`'s complement.
+        return not self.isdir
+
+    def __str__(self):
+        return f"Metadata: {self.path_info}"

dvc/tree/dvc.py

@@ -4,6 +4,7 @@ import os
 from dvc.exceptions import OutputNotFoundError
 from dvc.path_info import PathInfo
 
+from ._metadata import Metadata
 from .base import BaseTree, RemoteActionNotImplemented
 
 logger = logging.getLogger(__name__)
@@ -100,17 +101,19 @@ class DvcTree(BaseTree):  # pylint:disable=abstract-method
 
     def exists(self, path):  # pylint: disable=arguments-differ
         try:
-            self._find_outs(path, strict=False, recursive=True)
+            self.metadata(path)
             return True
         except OutputNotFoundError:
             return False
 
     def isdir(self, path):  # pylint: disable=arguments-differ
-        if not self.exists(path):
+        try:
+            meta = self.metadata(path)
+            return meta.isdir
+        except OutputNotFoundError:
             return False
 
-        path_info = PathInfo(os.path.abspath(path))
-        outs = self._find_outs(path, strict=False, recursive=True)
+    def check_isdir(self, path_info, outs):
         if len(outs) != 1:
             return True
 
@@ -129,11 +132,12 @@ class DvcTree(BaseTree):  # pylint:disable=abstract-method
             return True
 
     def isfile(self, path):  # pylint: disable=arguments-differ
-        if not self.exists(path):
+        try:
+            meta = self.metadata(path)
+            return meta.isfile
+        except OutputNotFoundError:
             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
@@ -186,23 +190,21 @@ class DvcTree(BaseTree):  # pylint:disable=abstract-method
         from pygtrie import Trie
 
         assert topdown
-
-        if not self.exists(top):
+        root = PathInfo(os.path.abspath(top))
+        try:
+            meta = self.metadata(root)
+        except OutputNotFoundError:
             if onerror is not None:
                 onerror(FileNotFoundError(top))
             return
 
-        if not self.isdir(top):
+        if not meta.isdir:
             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:
+        for out in meta.outs:
             trie[out.path_info.parts] = out
 
             if out.is_dir_checksum and root.isin_or_eq(out.path_info):
@@ -210,12 +212,14 @@ class DvcTree(BaseTree):  # pylint:disable=abstract-method
 
         yield from self._walk(root, trie, topdown=topdown, **kwargs)
 
-    def isdvc(self, path, **kwargs):
+    def isdvc(self, path, recursive=False, strict=True):
         try:
-            return len(self._find_outs(path, **kwargs)) == 1
+            meta = self.metadata(path)
         except OutputNotFoundError:
-            pass
-        return False
+            return False
+
+        recurse = recursive or not strict
+        return meta.output_exists if recurse else meta.is_output
 
     def isexec(self, path):  # pylint: disable=unused-argument
         return False
@@ -231,3 +235,11 @@ class DvcTree(BaseTree):  # pylint:disable=abstract-method
                 self._get_granular_checksum(path_info, out),
             )
         return out.tree.PARAM_CHECKSUM, out.checksum
+
+    def metadata(self, path_info):
+        path_info = PathInfo(os.path.abspath(path_info))
+        outs = self._find_outs(path_info, strict=False, recursive=True)
+
+        meta = Metadata(path_info=path_info, outs=outs)
+        meta.isdir = meta.isdir or self.check_isdir(meta.path_info, meta.outs)
+        return meta

dvc/tree/git.py

@@ -1,13 +1,12 @@
 import errno
 import io
 import os
-import stat
 
 from funcy import cached_property
 
 from dvc.exceptions import DvcException
 from dvc.tree.base import BaseTree
-from dvc.utils import relpath
+from dvc.utils import is_exec, relpath
 
 # see git-fast-import(1)
 GIT_MODE_DIR = 0o40000
@@ -197,7 +196,7 @@ class GitTree(BaseTree):  # pylint:disable=abstract-method
             return False
 
         mode = self.stat(path).st_mode
-        return mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+        return is_exec(mode)
 
     def stat(self, path):
         import git

dvc/tree/local.py

@@ -10,7 +10,7 @@ from dvc.exceptions import DvcException
 from dvc.path_info import PathInfo
 from dvc.scheme import Schemes
 from dvc.system import System
-from dvc.utils import file_md5, relpath, tmp_fname
+from dvc.utils import file_md5, is_exec, relpath, tmp_fname
 from dvc.utils.fs import (
     copy_fobj_to_file,
     copyfile,
@@ -157,7 +157,7 @@ class LocalTree(BaseTree):
 
     def isexec(self, path):
         mode = os.stat(path).st_mode
-        return mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+        return is_exec(mode)
 
     @staticmethod
     def stat(path):

dvc/tree/repo.py

@@ -1,6 +1,8 @@
 import logging
 import os
+import stat
 import threading
+from contextlib import suppress
 from itertools import takewhile
 from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union
 
@@ -10,9 +12,10 @@ from pygtrie import StringTrie
 from dvc.dvcfile import is_valid_filename
 from dvc.exceptions import OutputNotFoundError
 from dvc.path_info import PathInfo
-from dvc.utils import file_md5
+from dvc.utils import file_md5, is_exec
 from dvc.utils.fs import copy_fobj_to_file, makedirs
 
+from ._metadata import Metadata
 from .base import BaseTree
 from .dvc import DvcTree
 
@@ -345,3 +348,28 @@ class RepoTree(BaseTree):  # pylint:disable=abstract-method
     @property
     def hash_jobs(self):  # pylint: disable=invalid-overridden-method
         return self._main_repo.tree.hash_jobs
+
+    def metadata(self, path):
+        path_info = PathInfo(os.path.abspath(path))
+        tree, dvc_tree = self._get_tree_pair(path_info)
+
+        dvc_meta = None
+        if dvc_tree:
+            with suppress(OutputNotFoundError):
+                dvc_meta = dvc_tree.metadata(path_info)
+
+        stat_result = None
+        with suppress(FileNotFoundError):
+            stat_result = tree.stat(path_info)
+
+        if not stat_result and not dvc_meta:
+            raise FileNotFoundError
+
+        meta = dvc_meta or Metadata(path_info=path_info)
+
+        isdir = bool(stat_result) and stat.S_ISDIR(stat_result.st_mode)
+        meta.isdir = meta.isdir or isdir
+
+        if not dvc_meta:
+            meta.is_exec = bool(stat_result) and is_exec(stat_result.st_mode)
+        return meta

dvc/utils/__init__.py

@@ -6,6 +6,7 @@ import logging
 import math
 import os
 import re
+import stat
 import sys
 import time
 
@@ -415,3 +416,7 @@ def parse_target(target, default=None):
         logger.debug("Assuming file to be '%s'", path)
 
     return path, name
+
+
+def is_exec(mode):
+    return mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)

setup.py

@@ -74,6 +74,7 @@ install_requires = [
     "networkx>=2.1,<2.5",
     "pydot>=1.2.4",
     "speedcopy>=2.0.1; python_version < '3.8' and sys_platform == 'win32'",
+    "dataclasses; python_version < '3.7'",
     "flatten_json>=0.1.6,<0.1.8",
     "tabulate>=0.8.7",
     "pygtrie==2.3.2",

tests/unit/tree/test_repo.py

@@ -177,7 +177,7 @@ def test_isdvc(tmp_dir, dvc):
     assert not tree.isdvc("bar")
     assert tree.isdvc("dir")
     assert not tree.isdvc("dir/baz")
-    assert tree.isdvc("dir/baz", recursive=True, strict=False)
+    assert tree.isdvc("dir/baz", recursive=True)
 
 
 def make_subrepo(dir_, scm, config=None):

tests/unit/tree/test_repo_metadata.py

@@ -0,0 +1,195 @@
+import pytest
+
+from dvc.path_info import PathInfo
+from dvc.tree.repo import RepoTree
+
+
+@pytest.fixture(scope="module")
+def monkey_mod():
+    # using internal API, `monkeypatch` is a `function` scoped fixture
+    from _pytest.monkeypatch import MonkeyPatch
+
+    patch = MonkeyPatch()
+    yield patch
+    patch.undo()
+
+
+@pytest.fixture(scope="module")
+def temp_repo(tmp_path_factory, make_tmp_dir, monkey_mod):
+    path = tmp_path_factory.mktemp("temp-repo")
+    monkey_mod.chdir(path)
+    return make_tmp_dir(path, scm=True, dvc=True)
+
+
+@pytest.fixture(scope="module")
+def repo_tree(temp_repo):
+    fs_structure = {
+        "models": {  # mixed dvc + git directory
+            "train.py": "train dot py",
+            "test.py": "test dot py",
+        },
+        "README.md": "my little project",  # file
+        "src": {  # repo-only directory
+            "utils": {
+                "__init__.py": "",
+                "serve_model.py": "# this will serve a model `soon`",
+            }
+        },
+    }
+    dvc_structure = {
+        "data": {  # dvc only directory
+            "raw": {
+                "raw-1.csv": "one, dot, csv",
+                "raw-2.csv": "two, dot, csv",
+            },
+            "processed": {
+                "processed-1.csv": "1, dot, csv",
+                "processed-2.csv": "2, dot, csv",
+            },
+        },
+        "models/transform.pickle": "model model",  # file
+    }
+
+    temp_repo.scm_gen(fs_structure, commit="repo init")
+    temp_repo.dvc_gen(dvc_structure, commit="use dvc")
+
+    yield RepoTree(temp_repo.dvc, fetch=True)
+
+
+def test_metadata_not_existing(repo_tree):
+    path = PathInfo("path") / "that" / "does" / "not" / "exist"
+
+    with pytest.raises(FileNotFoundError):
+        repo_tree.metadata(path)
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "README.md",
+        "models/train.py",
+        "models/test.py",
+        PathInfo("src") / "utils" / "__init__.py",
+        PathInfo("src") / "utils" / "serve_model.py",
+    ],
+)
+def test_metadata_git_tracked_file(repo_tree, path):
+    root = PathInfo(repo_tree.root_dir)
+    meta = repo_tree.metadata(path)
+
+    assert meta.path_info == root / path
+    assert not meta.is_output
+    assert not meta.part_of_output
+    assert not meta.contains_outputs
+    assert not meta.is_dvc
+    assert not meta.output_exists
+    assert not meta.isdir
+    assert not meta.is_exec
+    assert meta.isfile
+    assert not meta.outs
+
+
+@pytest.mark.parametrize(
+    "path, outs, is_output",
+    [
+        (PathInfo("data") / "raw" / "raw-1.csv", [PathInfo("data")], False),
+        (PathInfo("data") / "raw" / "raw-2.csv", [PathInfo("data")], False),
+        (
+            PathInfo("data") / "processed" / "processed-1.csv",
+            [PathInfo("data")],
+            False,
+        ),
+        (
+            PathInfo("data") / "processed" / "processed-2.csv",
+            [PathInfo("data")],
+            False,
+        ),
+        (
+            "models/transform.pickle",
+            [PathInfo("models") / "transform.pickle"],
+            True,
+        ),
+    ],
+)
+def test_metadata_dvc_tracked_file(repo_tree, path, outs, is_output):
+    root = PathInfo(repo_tree.root_dir)
+    meta = repo_tree.metadata(path)
+
+    assert meta.path_info == root / path
+    assert meta.is_output == is_output
+    assert meta.part_of_output != is_output
+    assert not meta.contains_outputs
+    assert meta.is_dvc
+    assert meta.output_exists
+    assert not meta.isdir
+    assert not meta.is_exec
+    assert meta.isfile
+    assert {out.path_info for out in meta.outs} == {root / out for out in outs}
+
+
+@pytest.mark.parametrize("path", ["src", "src/utils"])
+def test_metadata_git_only_dirs(repo_tree, path):
+    root = PathInfo(repo_tree.root_dir)
+    meta = repo_tree.metadata(path)
+
+    assert meta.path_info == root / path
+    assert not meta.is_output
+    assert not meta.part_of_output
+    assert not meta.contains_outputs
+    assert not meta.is_dvc
+    assert not meta.output_exists
+    assert meta.isdir
+    assert meta.is_exec
+    assert not meta.isfile
+    assert not meta.outs
+
+
+@pytest.mark.parametrize(
+    "path, expected_outs",
+    [
+        (".", [PathInfo("data"), PathInfo("models") / "transform.pickle"]),
+        ("models", [PathInfo("models") / "transform.pickle"]),
+    ],
+)
+def test_metadata_git_dvc_mixed_dirs(repo_tree, path, expected_outs):
+    root = PathInfo(repo_tree.root_dir)
+    meta = repo_tree.metadata(root / path)
+
+    assert meta.path_info == root / path
+    assert not meta.is_output
+    assert not meta.part_of_output
+    assert meta.contains_outputs
+    assert not meta.is_dvc
+    assert meta.output_exists
+    assert meta.isdir
+    assert not meta.is_exec
+    assert not meta.isfile
+
+    assert {out.path_info for out in meta.outs} == {
+        root / out for out in expected_outs
+    }
+
+
+@pytest.mark.parametrize(
+    "path, is_output",
+    [
+        ("data", True),
+        ("data/raw", False),  # is inside output
+        ("data/processed", False),
+    ],
+)
+def test_metadata_dvc_only_dirs(repo_tree, path, is_output):
+    data = PathInfo(repo_tree.root_dir) / "data"
+    root = PathInfo(repo_tree.root_dir)
+    meta = repo_tree.metadata(root / path)
+
+    assert meta.path_info == root / path
+    assert meta.is_output == is_output
+    assert meta.part_of_output != is_output
+    assert not meta.contains_outputs
+    assert meta.is_dvc
+    assert meta.output_exists
+    assert meta.isdir
+    assert not meta.is_exec
+    assert not meta.isfile
+    assert {out.path_info for out in meta.outs} == {data}