diff --git a/dvc/data/tree.py b/dvc/data/tree.py index 25e29d812..c3f42b9d0 100644 --- a/dvc/data/tree.py +++ b/dvc/data/tree.py @@ -197,6 +197,27 @@ class Tree(HashFile): tree.digest() return tree + def ls(self, prefix=None): + kwargs = {} + if prefix: + kwargs["prefix"] = prefix + + meta, hash_info = self._trie.get(prefix, (None, None)) + if hash_info and hash_info.isdir and meta and not meta.obj: + raise FileNotFoundError + + ret = [] + + def node_factory(_, key, children, *args): + if key == prefix: + list(children) + else: + ret.append(key[-1]) + + self._trie.traverse(node_factory, **kwargs) + + return ret + def du(odb, tree): try: diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index cbe457139..e2b25585f 100644 --- a/dvc/fs/dvc.py +++ b/dvc/fs/dvc.py @@ -97,33 +97,35 @@ class DvcFileSystem(FileSystem): # pylint:disable=abstract-method except FileNotFoundError: return False + def ls(self, path, detail=True, **kwargs): + info = self.info(path) + if info["type"] != "directory": + return [info] if detail else [path] + + root_key = self._get_key(path) + try: + entries = [ + self.sep.join((path, name)) if path else name + for name in self.repo.index.tree.ls(prefix=root_key) + ] + except KeyError as exc: + raise FileNotFoundError from exc + + if not detail: + return entries + + return [self.info(epath) for epath in entries] + def _walk(self, root, topdown=True, **kwargs): - dirs = set() + dirs = [] files = [] - root_parts = self._get_key(root) - root_len = len(root_parts) - try: - for key, (meta, hash_info) in self.repo.index.tree.iteritems( - prefix=root_parts - ): # noqa: B301 - if hash_info and hash_info.isdir and meta and not meta.obj: - raise FileNotFoundError - - if key == root_parts: - continue - - if hash_info.isdir: - continue - - name = key[root_len] - if len(key) > root_len + 1: - dirs.add(name) - continue - + for entry in self.ls(root, detail=True): + name = self.path.name(entry["name"]) + if entry["type"] == "directory": + dirs.append(name) + else: files.append(name) - except KeyError: - pass assert topdown dirs = list(dirs) @@ -175,6 +177,7 @@ class DvcFileSystem(FileSystem): # pylint:disable=abstract-method "isexec": False, "isdvc": False, "outs": outs, + "name": path, } if len(outs) > 1 and outs[0][0] != key: