Repository: iterative/dvc
Pull request #10153 (2023-12-11 18:33:03)

[pre-commit.ci] pre-commit autoupdate
<!--pre-commit.ci start-->
updates:
- [github.com/astral-sh/ruff-pre-commit: v0.1.6 → v0.1.7](https://github.com/astral-sh/ruff-pre-commit/compare/v0.1.6...v0.1.7)
<!--pre-commit.ci end-->

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/iterative/dvc/pull/10153?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`77b451d`)](https://app.codecov.io/gh/iterative/dvc/commit/77b451d1184529e4b1df87bb0e25ffdec91d10ef?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) 90.63% compared to head [(`5aa825c`)](https://app.codecov.io/gh/iterative/dvc/pull/10153?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) 90.22%.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main   #10153      +/-   ##
==========================================
- Coverage   90.63%   90.22%   -0.41%     
==========================================
  Files         496      496              
  Lines       37751    37751              
  Branches     5490     5481       -9     
==========================================
- Hits        34215    34062     -153     
- Misses       2900     3007     +107     
- Partials      636      682      +46     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/iterative/dvc/pull/10153?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).


Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 98e909ea7..1def6d319 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -22,7 +22,7 @@ repos:
       - id: sort-simple-yaml
       - id: trailing-whitespace
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.1.6'
+    rev: 'v0.1.7'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
diff --git a/dvc/cli/parser.py b/dvc/cli/parser.py
index 31e68ad4d..eb0b456e8 100644
--- a/dvc/cli/parser.py
+++ b/dvc/cli/parser.py
@@ -19,6 +19,7 @@ from dvc.commands import (
     data_sync,
     destroy,
     diff,
+    du,
     experiments,
     freeze,
     gc,
@@ -95,6 +96,7 @@ COMMANDS = [
     data,
     artifacts,
     studio,
+    du,
 ]
 
 
diff --git a/dvc/commands/du.py b/dvc/commands/du.py
new file mode 100644
index 000000000..11f067c32
--- /dev/null
+++ b/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 total disk usage.",
+    )
+    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)
diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py
index be0df0ecc..8888999dd 100644
--- a/dvc/data_cloud.py
+++ b/dvc/data_cloud.py
@@ -3,7 +3,6 @@
 from typing import TYPE_CHECKING, Iterable, Optional, Set, Tuple
 
 from dvc.config import NoRemoteError, RemoteConfigError
-from dvc.fs.callbacks import Callback
 from dvc.log import logger
 from dvc.utils.objects import cached_property
 from dvc_data.hashfile.db import get_index
@@ -206,11 +205,13 @@ class DataCloud:
         jobs: Optional[int] = None,
         odb: "HashFileDB",
     ) -> "TransferResult":
+        from dvc.fs.callbacks import TqdmCallback
+
         if odb.hash_name == "md5-dos2unix":
             cache = self.repo.cache.legacy
         else:
             cache = self.repo.cache.local
-        with Callback.as_tqdm_callback(
+        with TqdmCallback(
             desc=f"Pushing to {odb.fs.unstrip_protocol(odb.path)}",
             unit="file",
         ) as cb:
@@ -265,11 +266,13 @@ class DataCloud:
         jobs: Optional[int] = None,
         odb: "HashFileDB",
     ) -> "TransferResult":
+        from dvc.fs.callbacks import TqdmCallback
+
         if odb.hash_name == "md5-dos2unix":
             cache = self.repo.cache.legacy
         else:
             cache = self.repo.cache.local
-        with Callback.as_tqdm_callback(
+        with TqdmCallback(
             desc=f"Fetching from {odb.fs.unstrip_protocol(odb.path)}",
             unit="file",
         ) as cb:
diff --git a/dvc/fs/__init__.py b/dvc/fs/__init__.py
index 1e8627efc..3404df217 100644
--- a/dvc/fs/__init__.py
+++ b/dvc/fs/__init__.py
@@ -27,7 +27,7 @@ from dvc_objects.fs.errors import (  # noqa: F401
 )
 from dvc_objects.fs.path import Path  # noqa: F401
 
-from .callbacks import Callback
+from .callbacks import Callback  # noqa: F401
 from .data import DataFileSystem  # noqa: F401
 from .dvc import DVCFileSystem
 from .git import GitFileSystem  # noqa: F401
@@ -51,7 +51,9 @@ def download(
 ) -> int:
     from dvc.scm import lfs_prefetch
 
-    with Callback.as_tqdm_callback(
+    from .callbacks import TqdmCallback
+
+    with TqdmCallback(
         desc=f"Downloading {fs.path.name(fs_path)}",
         unit="files",
     ) as cb:
diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py
index ccc2c944a..25926c45f 100644
--- a/dvc/fs/dvc.py
+++ b/dvc/fs/dvc.py
@@ -4,6 +4,7 @@ import ntpath
 import os
 import posixpath
 import threading
+from collections import deque
 from contextlib import ExitStack, suppress
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type, Union
 
@@ -60,6 +61,7 @@ def _merge_info(repo, key, fs_info, dvc_info):
     if fs_info:
         ret["type"] = fs_info["type"]
         ret["size"] = fs_info["size"]
+        ret["fs_info"] = fs_info
         isexec = False
         if fs_info["type"] == "file":
             isexec = utils.is_exec(fs_info["mode"])
@@ -421,6 +423,45 @@ class _DVCFileSystem(AbstractFileSystem):
         dvc_path = _get_dvc_path(dvc_fs, subkey)
         return dvc_fs.get_file(dvc_path, lpath, **kwargs)
 
+    def du(self, path, total=True, maxdepth=None, withdirs=False, **kwargs):
+        if maxdepth is not None:
+            raise NotImplementedError
+
+        sizes = {}
+        dus = {}
+        todo = deque([self.info(path)])
+        while todo:
+            info = todo.popleft()
+            isdir = info["type"] == "directory"
+            size = info["size"] or 0
+            name = info["name"]
+
+            if not isdir:
+                sizes[name] = size
+                continue
+
+            dvc_info = info.get("dvc_info") or {}
+            fs_info = info.get("fs_info")
+            entry = dvc_info.get("entry")
+            if (
+                dvc_info
+                and not fs_info
+                and entry is not None
+                and entry.size is not None
+            ):
+                dus[name] = entry.size
+                continue
+
+            if withdirs:
+                sizes[name] = size
+
+            todo.extend(self.ls(info["name"], detail=True))
+
+        if total:
+            return sum(sizes.values()) + sum(dus.values())
+
+        return sizes
+
     def close(self):
         self._repo_stack.close()
 
diff --git a/dvc/output.py b/dvc/output.py
index a7278ce69..decc97e05 100644
--- a/dvc/output.py
+++ b/dvc/output.py
@@ -786,7 +786,7 @@ class Output:
                     self.hash_name,
                     ignore=self.dvcignore,
                 )
-                with Callback.as_tqdm_callback(
+                with TqdmCallback(
                     desc=f"Committing {self} to cache",
                     unit="file",
                 ) as cb:
@@ -826,7 +826,7 @@ class Output:
         save_obj = obj.filter(prefix)
         assert isinstance(save_obj, Tree)
         checkout_obj = save_obj.get_obj(self.cache, prefix)
-        with Callback.as_tqdm_callback(
+        with TqdmCallback(
             desc=f"Committing {self} to cache",
             unit="file",
         ) as cb:
@@ -1059,7 +1059,7 @@ class Output:
             upload=upload,
             no_progress_bar=no_progress_bar,
         )
-        with Callback.as_tqdm_callback(
+        with TqdmCallback(
             desc=f"Transferring to {odb.fs.unstrip_protocol(odb.path)}",
             unit="file",
         ) as cb:
@@ -1092,7 +1092,10 @@ class Output:
 
     def unprotect(self):
         if self.exists and self.use_cache:
-            self.cache.unprotect(self.fs_path)
+            with TqdmCallback(
+                size=self.meta.nfiles or -1, desc=f"Unprotecting {self}"
+            ) as callback:
+                self.cache.unprotect(self.fs_path, callback=callback)
 
     def get_dir_cache(self, **kwargs) -> Optional["Tree"]:
         if not self.is_dir_checksum:
@@ -1435,7 +1438,7 @@ class Output:
 
         assert staging
         assert obj.hash_info
-        with Callback.as_tqdm_callback(
+        with TqdmCallback(
             desc=f"Adding {self} to cache",
             unit="file",
         ) as cb:
diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py
index 804d5a25b..2c6f965bf 100644
--- a/dvc/repo/__init__.py
+++ b/dvc/repo/__init__.py
@@ -70,6 +70,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 as _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]
@@ -93,6 +94,7 @@ class Repo:
     from .cache import check_missing as cache_check_missing  # type: ignore[misc]
     from .data import status as data_status  # type: ignore[misc]
 
+    du = staticmethod(_du)
     ls = staticmethod(_ls)
     ls_url = staticmethod(_ls_url)
     get = staticmethod(_get)
@@ -447,7 +449,10 @@ class Repo:
         return init(root_dir=root_dir, no_scm=no_scm, force=force, subdir=subdir)
 
     def unprotect(self, target):
-        return self.cache.repo.unprotect(target)
+        from dvc.fs.callbacks import TqdmCallback
+
+        with TqdmCallback(desc=f"Unprotecting {target}") as callback:
+            return self.cache.repo.unprotect(target, callback=callback)
 
     def _ignore(self):
         flist = [self.config.files["local"]]
diff --git a/dvc/repo/artifacts.py b/dvc/repo/artifacts.py
index 34ee01f24..87a62ec7e 100644
--- a/dvc/repo/artifacts.py
+++ b/dvc/repo/artifacts.py
@@ -210,7 +210,8 @@ class Artifacts:
     ) -> Tuple[int, str]:
         from dvc_studio_client.model_registry import get_download_uris
 
-        from dvc.fs import Callback, HTTPFileSystem, generic, localfs
+        from dvc.fs import HTTPFileSystem, generic, localfs
+        from dvc.fs.callbacks import TqdmCallback
 
         logger.debug("Trying to download artifact '%s' via studio", name)
         out = out or os.getcwd()
@@ -244,7 +245,7 @@ class Artifacts:
             ) from exc
         fs = HTTPFileSystem()
         jobs = jobs or fs.jobs
-        with Callback.as_tqdm_callback(
+        with TqdmCallback(
             desc=f"Downloading '{name}' from '{repo_url}'",
             unit="files",
         ) as cb:
diff --git a/dvc/repo/du.py b/dvc/repo/du.py
new file mode 100644
index 000000000..f51148088
--- /dev/null
+++ b/dvc/repo/du.py
@@ -0,0 +1,42 @@
+from typing import Any, Dict, Optional, Union
+
+
+def du(
+    url: str,
+    path: Optional[str] = None,
+    rev: Optional[str] = None,
+    summarize: bool = False,
+    config: Union[None, Dict[str, Any], 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)
+        ]
+        ret.append((path, sum(entry[1] for entry in ret)))
+        return ret
diff --git a/dvc/repo/imports.py b/dvc/repo/imports.py
index 8e377fc97..c373d7ace 100644
--- a/dvc/repo/imports.py
+++ b/dvc/repo/imports.py
@@ -100,10 +100,10 @@ def save_imports(
     Returns:
         Objects which were downloaded from source location.
     """
+    from dvc.fs.callbacks import TqdmCallback
     from dvc.stage.exceptions import DataSourceChanged
     from dvc_data.index import md5, save
     from dvc_data.index.checkout import apply, compare
-    from dvc_objects.fs.callbacks import Callback
 
     downloaded: Set["HashInfo"] = set()
 
@@ -122,7 +122,7 @@ def save_imports(
             if not cache.fs.exists(cache.path):
                 os.makedirs(cache.path)
             with TemporaryDirectory(dir=cache.path) as tmpdir:
-                with Callback.as_tqdm_callback(
+                with TqdmCallback(
                     desc="Downloading imports from source",
                     unit="files",
                 ) as cb:
diff --git a/dvc/repo/worktree.py b/dvc/repo/worktree.py
index 8b8d4a95a..023575498 100644
--- a/dvc/repo/worktree.py
+++ b/dvc/repo/worktree.py
@@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Tuple, Union
 
 from funcy import first
 
-from dvc.fs.callbacks import Callback
 from dvc.log import logger
 from dvc.stage.exceptions import StageUpdateError
 
@@ -247,11 +246,12 @@ def _fetch_out_changes(
     remote_index: Union["DataIndex", "DataIndexView"],
     remote: "Remote",
 ):
+    from dvc.fs.callbacks import TqdmCallback
     from dvc_data.index.checkout import apply, compare
 
     old, new = _get_diff_indexes(out, local_index, remote_index)
 
-    with Callback.as_tqdm_callback(
+    with TqdmCallback(
         unit="entry",
         desc="Comparing indexes",
     ) as cb:
@@ -265,9 +265,7 @@ def _fetch_out_changes(
         )
 
     total = len(new)
-    with Callback.as_tqdm_callback(
-        unit="file", desc=f"Updating '{out}'", disable=total == 0
-    ) as cb:
+    with TqdmCallback(unit="file", desc=f"Updating '{out}'", disable=total == 0) as cb:
         cb.set_size(total)
         apply(
             diff,
diff --git a/dvc/stage/cache.py b/dvc/stage/cache.py
index f8807191c..4e8501eb1 100644
--- a/dvc/stage/cache.py
+++ b/dvc/stage/cache.py
@@ -231,7 +231,7 @@ class StageCache:
 
     def transfer(self, from_odb, to_odb):
         from dvc.fs import HTTPFileSystem, LocalFileSystem
-        from dvc.fs.callbacks import Callback
+        from dvc.fs.callbacks import TqdmCallback
 
         from_fs = from_odb.fs
         to_fs = to_odb.fs
@@ -266,7 +266,7 @@ class StageCache:
 
             src_name = from_fs.path.name(src)
             parent_name = from_fs.path.name(from_fs.path.parent(src))
-            with Callback.as_tqdm_callback(
+            with TqdmCallback(
                 desc=src_name,
                 bytes=True,
             ) as cb:
diff --git a/dvc/utils/plots.py b/dvc/utils/plots.py
index 6e85bea47..ee7f988c1 100644
--- a/dvc/utils/plots.py
+++ b/dvc/utils/plots.py
@@ -9,7 +9,7 @@ def get_plot_id(config_plot_id: str, config_file_path: str = ""):
 
 
 def group_definitions_by_id(
-    definitions: Dict[str, Dict]
+    definitions: Dict[str, Dict],
 ) -> Dict[str, Tuple[str, Dict]]:
     """
     Format ID and extracts plot_definition for each plot.
diff --git a/pyproject.toml b/pyproject.toml
index b6fb76835..29a29f6f9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,7 +35,7 @@ dependencies = [
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=2.22.6,<2.23.0",
+    "dvc-data>=2.24,<2.25",
     "dvc-http>=2.29.0",
     "dvc-render>=1.0.0,<2",
     "dvc-studio-client>=0.17.1,<1",


Test patch:
diff --git a/tests/func/test_du.py b/tests/func/test_du.py
new file mode 100644
index 000000000..9496f8ab5
--- /dev/null
+++ b/tests/func/test_du.py
@@ -0,0 +1,49 @@
+import os
+
+
+def test_du(tmp_dir, dvc):
+    tmp_dir.gen(
+        {
+            "file": b"file",
+            "dvcfile": b"dvcfile",
+            "dir": {
+                "dirfile": b"dirfile",
+                "subdir": {
+                    "subdirfile": b"subdirfile",
+                },
+                "dvcsubdir": {
+                    "dvcsubdirfile": b"dvcsubdirfile",
+                },
+            },
+        }
+    )
+
+    dvc.add("dvcfile")
+    dvc.add(os.path.join("dir", "dvcsubdir"))
+
+    assert dvc.du(".", "file") == [("file", 4)]
+    assert dvc.du(".", "dvcfile") == [("dvcfile", 7)]
+    assert set(dvc.du(".", "dir/subdir")) == {
+        ("dir/subdir/subdirfile", 10),
+        ("dir/subdir", 10),
+    }
+    assert dvc.du(".", "dir/subdir", summarize=True) == [("dir/subdir", 10)]
+    assert set(dvc.du(".", "dir/dvcsubdir")) == {
+        ("dir/dvcsubdir/dvcsubdirfile", 13),
+        ("dir/dvcsubdir", 13),
+    }
+    assert dvc.du(".", "dir/dvcsubdir", summarize=True) == [("dir/dvcsubdir", 13)]
+    assert set(dvc.du(".", "dir")) == {
+        ("dir/dvcsubdir", 13),
+        ("dir/subdir", 10),
+        ("dir/dirfile", 7),
+        ("dir", 30),
+    }
+    assert dvc.du(".", "dir", summarize=True) == [("dir", 30)]
+    assert set(dvc.du(".", "/")) == {
+        ("/dvcfile", 7),
+        ("/dir", 30),
+        ("/file", 4),
+        ("/", 41),
+    }
+    assert dvc.du(".", "/", summarize=True) == [("/", 41)]
diff --git a/tests/unit/command/test_du.py b/tests/unit/command/test_du.py
new file mode 100644
index 000000000..6ad99494e
--- /dev/null
+++ b/tests/unit/command/test_du.py
@@ -0,0 +1,21 @@
+from dvc.cli import parse_args
+from dvc.commands.du import CmdDU
+
+
+def test_du(mocker):
+    cli_args = parse_args(["du", "myurl", "mypath", "--summarize", "--rev", "myrev"])
+    assert cli_args.func == CmdDU
+
+    cmd = cli_args.func(cli_args)
+    mock_du = mocker.patch("dvc.repo.Repo.du")
+
+    assert cmd.run() == 0
+    mock_du.assert_called_once_with(
+        "myurl",
+        "mypath",
+        rev="myrev",
+        summarize=True,
+        config=None,
+        remote=None,
+        remote_config=None,
+    )