ls: implement --tree to show as a tree and --level to limit depth for recursion

Author: skshetry

dvc/commands/ls/__init__.py

@@ -1,3 +1,5 @@
+from typing import Callable
+
 from dvc.cli import completion, formatter
 from dvc.cli.command import CmdBaseNoRepo
 from dvc.cli.utils import DictAction, append_doc_link
@@ -9,7 +11,18 @@
 logger = logger.getChild(__name__)
 
 
-def _format_entry(entry, fmt, with_size=True, with_md5=False):
+def _get_formatter(with_color: bool = False) -> Callable[[dict], str]:
+    def fmt(entry: dict) -> str:
+        return entry["path"]
+
+    if with_color:
+        ls_colors = LsColors()
+        return ls_colors.format
+
+    return fmt
+
+
+def _format_entry(entry, name, with_size=True, with_hash=False):
     from dvc.utils.humanize import naturalsize
 
     ret = []
@@ -20,60 +33,156 @@ def _format_entry(entry, fmt, with_size=True, with_md5=False):
         else:
             size = naturalsize(size)
         ret.append(size)
-    if with_md5:
+    if with_hash:
         md5 = entry.get("md5", "")
         ret.append(md5)
-    ret.append(fmt(entry))
+    ret.append(name)
     return ret
 
 
-def show_entries(entries, with_color=False, with_size=False, with_md5=False):
-    if with_color:
-        ls_colors = LsColors()
-        fmt = ls_colors.format
-    else:
-
-        def fmt(entry):
-            return entry["path"]
-
-    if with_size or with_md5:
+def show_entries(entries, with_color=False, with_size=False, with_hash=False):
+    fmt = _get_formatter(with_color)
+    if with_size or with_hash:
+        colalign = ("right",) if with_size else None
         ui.table(
             [
-                _format_entry(entry, fmt, with_size=with_size, with_md5=with_md5)
+                _format_entry(
+                    entry,
+                    fmt(entry),
+                    with_size=with_size,
+                    with_hash=with_hash,
+                )
                 for entry in entries
-            ]
+            ],
+            colalign=colalign,
         )
         return
 
     # NOTE: this is faster than ui.table for very large number of entries
     ui.write("\n".join(fmt(entry) for entry in entries))
 
 
+class TreePart:
+    Edge = "├── "
+    Line = "│   "
+    Corner = "└── "
+    Blank = "    "
+
+
+def _build_tree_structure(
+    entries, with_color=False, with_size=False, with_hash=False, _depth=0, _prefix=""
+):
+    rows = []
+    fmt = _get_formatter(with_color)
+
+    num_entries = len(entries)
+    for i, (name, entry) in enumerate(entries.items()):
+        entry["path"] = name
+        is_last = i >= num_entries - 1
+        tree_part = ""
+        if _depth > 0:
+            tree_part = TreePart.Corner if is_last else TreePart.Edge
+
+        row = _format_entry(
+            entry,
+            _prefix + tree_part + fmt(entry),
+            with_size=with_size,
+            with_hash=with_hash,
+        )
+        rows.append(row)
+
+        if contents := entry.get("contents"):
+            new_prefix = _prefix
+            if _depth > 0:
+                new_prefix += TreePart.Blank if is_last else TreePart.Line
+            new_rows = _build_tree_structure(
+                contents,
+                with_color=with_color,
+                with_size=with_size,
+                with_hash=with_hash,
+                _depth=_depth + 1,
+                _prefix=new_prefix,
+            )
+            rows.extend(new_rows)
+
+    return rows
+
+
+def show_tree(entries, with_color=False, with_size=False, with_hash=False):
+    import tabulate
+
+    rows = _build_tree_structure(
+        entries,
+        with_color=with_color,
+        with_size=with_size,
+        with_hash=with_hash,
+    )
+
+    colalign = ("right",) if with_size else None
+
+    _orig = tabulate.PRESERVE_WHITESPACE
+    tabulate.PRESERVE_WHITESPACE = True
+    try:
+        ui.table(rows, colalign=colalign)
+    finally:
+        tabulate.PRESERVE_WHITESPACE = _orig
+
+
 class CmdList(CmdBaseNoRepo):
-    def run(self):
+    def _show_tree(self):
+        from dvc.repo.ls import ls_tree
+
+        entries = ls_tree(
+            self.args.url,
+            self.args.path,
+            rev=self.args.rev,
+            dvc_only=self.args.dvc_only,
+            config=self.args.config,
+            remote=self.args.remote,
+            remote_config=self.args.remote_config,
+            maxdepth=self.args.level,
+        )
+        show_tree(
+            entries,
+            with_color=True,
+            with_size=self.args.size,
+            with_hash=self.args.show_hash,
+        )
+        return 0
+
+    def _show_list(self):
         from dvc.repo import Repo
 
-        try:
-            entries = Repo.ls(
-                self.args.url,
-                self.args.path,
-                rev=self.args.rev,
-                recursive=self.args.recursive,
-                dvc_only=self.args.dvc_only,
-                config=self.args.config,
-                remote=self.args.remote,
-                remote_config=self.args.remote_config,
+        entries = Repo.ls(
+            self.args.url,
+            self.args.path,
+            rev=self.args.rev,
+            recursive=self.args.recursive,
+            dvc_only=self.args.dvc_only,
+            config=self.args.config,
+            remote=self.args.remote,
+            remote_config=self.args.remote_config,
+            maxdepth=self.args.level,
+        )
+        if self.args.json:
+            ui.write_json(entries)
+        elif entries:
+            show_entries(
+                entries,
+                with_color=True,
+                with_size=self.args.size,
+                with_hash=self.args.show_hash,
             )
-            if self.args.json:
-                ui.write_json(entries)
-            elif entries:
-                show_entries(
-                    entries,
-                    with_color=True,
-                    with_size=self.args.size,
-                    with_md5=self.args.show_hash,
-                )
-            return 0
+        return 0
+
+    def run(self):
+        if self.args.tree and self.args.json:
+            raise DvcException("Cannot use --tree and --json options together.")
+
+        try:
+            if self.args.tree:
+                return self._show_tree()
+            return self._show_list()
         except FileNotFoundError:
             logger.exception("")
             return 1
@@ -102,6 +211,19 @@ def add_parser(subparsers, parent_parser):
         action="store_true",
         help="Recursively list files.",
     )
+    list_parser.add_argument(
+        "-T",
+        "--tree",
+        action="store_true",
+        help="Recurse into directories as a tree.",
+    )
+    list_parser.add_argument(
+        "-L",
+        "--level",
+        metavar="depth",
+        type=int,
+        help="Limit the depth of recursion.",
+    )
     list_parser.add_argument(
         "--dvc-only", action="store_true", help="Show only DVC outputs."
     )

dvc/commands/ls/ls_colors.py

@@ -32,7 +32,9 @@ def format(self, entry):
         if entry.get("isexec", False):
             return self._format(text, code="ex")
 
-        _, ext = os.path.splitext(text)
+        stem, ext = os.path.splitext(text)
+        if not ext and stem.startswith("."):
+            ext = stem
         return self._format(text, ext=ext)
 
     def _format(self, text, code=None, ext=None):