cli: introduce basic dvc du

Author: efiop

dvc/cli/parser.py

@@ -20,6 +20,7 @@
     data_sync,
     destroy,
     diff,
+    du,
     experiments,
     freeze,
     gc,
@@ -91,6 +92,7 @@
     check_ignore,
     data,
     artifacts,
+    du,
 ]
 
 

dvc/commands/du.py

@@ -0,0 +1,80 @@
+import argparse
+import logging
+
+from dvc.cli import completion
+from dvc.cli.command import CmdBaseNoRepo
+from dvc.cli.utils import DictAction, append_doc_link
+from dvc.ui import ui
+
+logger = logging.getLogger(__name__)
+
+
+class CmdDU(CmdBaseNoRepo):
+    def run(self):
+        from dvc.repo import Repo
+        from dvc.utils.humanize import naturalsize
+
+        entries = Repo.du(
+            self.args.url,
+            self.args.path,
+            rev=self.args.rev,
+            summarize=self.args.summarize,
+            config=self.args.config,
+            remote=self.args.remote,
+            remote_config=self.args.remote_config,
+        )
+        ui.table([(naturalsize(size), path) for path, size in entries])
+        return 0
+
+
+def add_parser(subparsers, parent_parser):
+    DU_HELP = "Show disk usage."
+    du_parser = subparsers.add_parser(
+        "du",
+        parents=[parent_parser],
+        description=append_doc_link(DU_HELP, "du"),
+        help=DU_HELP,
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+    du_parser.add_argument("url", help="Location of DVC repository")
+    du_parser.add_argument(
+        "--rev",
+        nargs="?",
+        help="Git revision (e.g. SHA, branch, tag)",
+        metavar="<commit>",
+    )
+    du_parser.add_argument(
+        "-s",
+        "--summarize",
+        action="store_true",
+        help="Show only total du.",
+    )
+    du_parser.add_argument(
+        "--config",
+        type=str,
+        help=(
+            "Path to a config file that will be merged with the config "
+            "in the target repository."
+        ),
+    )
+    du_parser.add_argument(
+        "--remote",
+        type=str,
+        help="Remote name to set as a default in the target repository.",
+    )
+    du_parser.add_argument(
+        "--remote-config",
+        type=str,
+        nargs="*",
+        action=DictAction,
+        help=(
+            "Remote config options to merge with a remote's config (default or one "
+            "specified by '--remote') in the target repository."
+        ),
+    )
+    du_parser.add_argument(
+        "path",
+        nargs="?",
+        help="Path to directory within the repository",
+    ).complete = completion.DIR
+    du_parser.set_defaults(func=CmdDU)

dvc/repo/__init__.py

@@ -71,6 +71,7 @@ class Repo:
     from dvc.repo.commit import commit  # type: ignore[misc]
     from dvc.repo.destroy import destroy  # type: ignore[misc]
     from dvc.repo.diff import diff  # type: ignore[misc]
+    from dvc.repo.du import du  # type: ignore[misc]
     from dvc.repo.fetch import fetch  # type: ignore[misc]
     from dvc.repo.freeze import freeze, unfreeze  # type: ignore[misc]
     from dvc.repo.gc import gc  # type: ignore[misc]

dvc/repo/du.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+
+def du(
+    url: str,
+    path: Optional[str] = None,
+    rev: Optional[str] = None,
+    summarize: bool = False,
+    config: Optional[str] = None,
+    remote: Optional[str] = None,
+    remote_config: Optional[dict] = None,
+):
+    from dvc.config import Config
+
+    from . import Repo
+
+    if config and not isinstance(config, dict):
+        config_dict = Config.load_file(config)
+    else:
+        config_dict = None
+
+    with Repo.open(
+        url,
+        rev=rev,
+        subrepos=True,
+        uninitialized=True,
+        config=config_dict,
+        remote=remote,
+        remote_config=remote_config,
+    ) as repo:
+        path = path or ""
+
+        fs = repo.dvcfs
+
+        if summarize or not fs.isdir(path):
+            return [(path, fs.du(path, total=True))]
+
+        ret = [
+            (entry_path, fs.du(entry_path, total=True))
+            for entry_path in fs.ls(path)  # pylint: disable=not-an-iterable
+        ]
+        ret.append((path, sum(entry[1] for entry in ret)))
+        return ret