cleaning up add, progress bars and status (#6224)

* cleanup code

* fix progress bar updating message even when index > total

* cleanup

dvc/repo/add.py

@@ -1,9 +1,13 @@
 import logging
 import os
-from typing import TYPE_CHECKING
+from contextlib import contextmanager
+from itertools import tee
+from typing import TYPE_CHECKING, Any, Iterator, List
 
 import colorama
 
+from dvc.ui import ui
+
 from ..exceptions import (
     CacheLinkError,
     InvalidArgumentError,
@@ -11,19 +15,28 @@ from ..exceptions import (
     OverlappingOutputPathsError,
     RecursiveAddingWhileUsingFilename,
 )
-from ..progress import Tqdm
 from ..repo.scm_context import scm_context
 from ..utils import LARGE_DIR_SIZE, glob_targets, resolve_output, resolve_paths
 from ..utils.collections import ensure_list, validate
 from . import locked
 
 if TYPE_CHECKING:
+    from dvc.repo import Repo
+    from dvc.stage import Stage
     from dvc.types import TargetType
 
-
+Stages = List["Stage"]
 logger = logging.getLogger(__name__)
 
 
+OVERLAPPING_OUTPUT_FMT = (
+    "Cannot add '{out}', because it is overlapping with other "
+    "DVC tracked output: '{parent}'.\n"
+    "To include '{out}' in '{parent}', run "
+    "'dvc commit {parent_stage}'"
+)
+
+
 def check_recursive_and_fname(args):
     if args.recursive and args.fname:
         raise RecursiveAddingWhileUsingFilename()
@@ -61,6 +74,61 @@ def check_arg_combinations(args):
         raise InvalidArgumentError(message.format(option=invalid_opt))
 
 
+@contextmanager
+def translate_graph_error(stages: Stages) -> Iterator[None]:
+    try:
+        yield
+    except OverlappingOutputPathsError as exc:
+        msg = OVERLAPPING_OUTPUT_FMT.format(
+            out=exc.overlapping_out.path_info,
+            parent=exc.parent.path_info,
+            parent_stage=exc.parent.stage.addressing,
+        )
+        raise OverlappingOutputPathsError(exc.parent, exc.overlapping_out, msg)
+    except OutputDuplicationError as exc:
+        raise OutputDuplicationError(
+            exc.output, list(set(exc.stages) - set(stages))
+        )
+
+
+def progress_iter(stages: Stages) -> Iterator["Stage"]:
+    total = len(stages)
+    desc = "Adding..."
+    with ui.progress(total=total, desc=desc, unit="file", leave=True) as pbar:
+        if total == 1:
+            pbar.bar_format = desc
+            pbar.refresh()
+
+        for stage in stages:
+            if total > 1:
+                pbar.update_msg(f"{stage.outs[0]}")
+            yield stage
+            if total == 1:  # restore bar format for stats
+                # pylint: disable=no-member
+                pbar.bar_format = pbar.BAR_FMT_DEFAULT
+
+
+LINK_FAILURE_MESSAGE = (
+    "\nSome targets could not be linked from cache to workspace.\n{}\n"
+    "To re-link these targets, reconfigure cache types and then run:\n"
+    "\n\tdvc checkout {}"
+)
+
+
+@contextmanager
+def warn_link_failures() -> Iterator[List[str]]:
+    link_failures: List[str] = []
+    try:
+        yield link_failures
+    finally:
+        if link_failures:
+            msg = LINK_FAILURE_MESSAGE.format(
+                CacheLinkError.SUPPORT_LINK,
+                " ".join(link_failures),
+            )
+            ui.error_write(msg)
+
+
 VALIDATORS = (
     check_recursive_and_fname,
     transform_targets,
@@ -68,226 +136,126 @@ VALIDATORS = (
 )
 
 
+@validate(*VALIDATORS)
 @locked
 @scm_context
-@validate(*VALIDATORS)
 def add(  # noqa: C901
-    repo,
+    repo: "Repo",
     targets: "TargetType",
-    recursive=False,
-    no_commit=False,
-    fname=None,
-    to_remote=False,
-    **kwargs,
+    recursive: bool = False,
+    no_commit: bool = False,
+    fname: str = None,
+    to_remote: bool = False,
+    **kwargs: Any,
 ):
-    link_failures = []
-    stages_list = []
-    num_targets = len(targets)
-    to_cache = kwargs.get("out") and not to_remote
-    with Tqdm(total=num_targets, desc="Add", unit="file", leave=True) as pbar:
-        if num_targets == 1:
-            # clear unneeded top-level progress bar for single target
-            pbar.bar_format = "Adding..."
-            pbar.refresh()
-        for target in targets:
-            sub_targets = _find_all_targets(repo, target, recursive)
-            pbar.total += len(sub_targets) - 1
+    to_cache = bool(kwargs.get("out")) and not to_remote
+    transfer = to_remote or to_cache
 
-            if os.path.isdir(target) and len(sub_targets) > LARGE_DIR_SIZE:
-                logger.warning(
-                    "You are adding a large directory '{target}' recursively,"
-                    " consider tracking it as a whole instead.\n"
-                    "{purple}HINT:{nc} Remove the generated DVC file and then"
-                    " run `{cyan}dvc add {target}{nc}`".format(
-                        purple=colorama.Fore.MAGENTA,
-                        cyan=colorama.Fore.CYAN,
-                        nc=colorama.Style.RESET_ALL,
-                        target=target,
-                    )
+    glob = kwargs.get("glob", False)
+    add_targets = collect_targets(repo, targets, recursive, glob)
+    # pass one for creating stages, other one is used for iterating here
+    add_targets, sources = tee(add_targets)
+
+    # collect targets and build stages as we go
+    desc = "Collecting targets"
+    stages_it = create_stages(repo, add_targets, fname, transfer, **kwargs)
+    stages = list(ui.progress(stages_it, desc=desc, unit="file"))
+
+    msg = "Collecting stages from the workspace"
+    with translate_graph_error(stages), ui.status(msg) as status:
+        # remove existing stages that are to-be replaced with these
+        # new stages for the graph checks.
+        old_stages = set(repo.stages) - set(stages)
+        status.update("Checking graph")
+        repo.check_modified_graph(stages, list(old_stages))
+
+    odb = None
+    if to_remote:
+        remote = repo.cloud.get_remote(kwargs.get("remote"), "add")
+        odb = remote.odb
+
+    with warn_link_failures() as link_failures:
+        for stage, source in zip(progress_iter(stages), sources):
+            if to_remote or to_cache:
+                stage.transfer(source, to_remote=to_remote, odb=odb, **kwargs)
+            else:
+                try:
+                    stage.save()
+                    if not no_commit:
+                        stage.commit()
+                except CacheLinkError:
+                    link_failures.append(str(stage.relpath))
+            stage.dump()
+    return stages
+
+
+LARGE_DIR_RECURSIVE_ADD_WARNING = (
+    "You are adding a large directory '{target}' recursively.\n"
+    "Consider tracking it as a whole instead with "
+    "`{cyan}dvc add {target}{nc}`."
+)
+
+
+def collect_targets(
+    repo: "Repo",
+    targets: "TargetType",
+    recursive: bool = False,
+    glob: bool = False,
+) -> Iterator[str]:
+    for target in glob_targets(ensure_list(targets), glob=glob):
+        expanded_targets = _find_all_targets(repo, target, recursive=recursive)
+        for index, path in enumerate(expanded_targets):
+            if index == LARGE_DIR_SIZE:
+                msg = LARGE_DIR_RECURSIVE_ADD_WARNING.format(
+                    cyan=colorama.Fore.CYAN,
+                    nc=colorama.Style.RESET_ALL,
+                    target=target,
                 )
-
-            stages = _create_stages(
-                repo,
-                sub_targets,
-                fname,
-                pbar=pbar,
-                transfer=to_remote or to_cache,
-                **kwargs,
-            )
-
-            # remove existing stages that are to-be replaced with these
-            # new stages for the graph checks.
-            old_stages = set(repo.stages) - set(stages)
-            try:
-                repo.check_modified_graph(stages, list(old_stages))
-            except OverlappingOutputPathsError as exc:
-                msg = (
-                    "Cannot add '{out}', because it is overlapping with other "
-                    "DVC tracked output: '{parent}'.\n"
-                    "To include '{out}' in '{parent}', run "
-                    "'dvc commit {parent_stage}'"
-                ).format(
-                    out=exc.overlapping_out.path_info,
-                    parent=exc.parent.path_info,
-                    parent_stage=exc.parent.stage.addressing,
-                )
-                raise OverlappingOutputPathsError(
-                    exc.parent, exc.overlapping_out, msg
-                )
-            except OutputDuplicationError as exc:
-                raise OutputDuplicationError(
-                    exc.output, list(set(exc.stages) - set(stages))
-                )
-
-            link_failures.extend(
-                _process_stages(
-                    repo,
-                    sub_targets,
-                    stages,
-                    no_commit,
-                    pbar,
-                    to_remote,
-                    to_cache,
-                    **kwargs,
-                )
-            )
-            stages_list += stages
-
-        if num_targets == 1:  # restore bar format for stats
-            pbar.bar_format = pbar.BAR_FMT_DEFAULT
-
-    if link_failures:
-        msg = (
-            "Some targets could not be linked from cache to workspace.\n{}\n"
-            "To re-link these targets, reconfigure cache types and then run:\n"
-            "\n\tdvc checkout {}"
-        ).format(
-            CacheLinkError.SUPPORT_LINK,
-            " ".join([str(stage.relpath) for stage in link_failures]),
-        )
-        logger.warning(msg)
-
-    return stages_list
+                ui.error_write(msg)
+            yield path
 
 
-def _process_stages(
-    repo, sub_targets, stages, no_commit, pbar, to_remote, to_cache, **kwargs
-):
-    link_failures = []
-    from dvc.dvcfile import Dvcfile
-
-    from ..output import OutputDoesNotExistError
-
-    if to_remote or to_cache:
-        # Already verified in the add()
-        (stage,) = stages
-        (target,) = sub_targets
-        (out,) = stage.outs
-
-        odb = None
-        if to_remote:
-            remote = repo.cloud.get_remote(kwargs.get("remote"), "add")
-            odb = remote.odb
-
-        out.transfer(target, odb=odb, jobs=kwargs.get("jobs"))
-
-        if to_cache:
-            out.checkout()
-
-        Dvcfile(repo, stage.path).dump(stage)
-        return link_failures
-
-    with Tqdm(
-        total=len(stages),
-        desc="Processing",
-        unit="file",
-        disable=len(stages) == 1,
-    ) as pbar_stages:
-        for stage in stages:
-            try:
-                stage.save()
-            except OutputDoesNotExistError:
-                pbar.n -= 1
-                raise
-
-            try:
-                if not no_commit:
-                    stage.commit()
-            except CacheLinkError:
-                link_failures.append(stage)
-
-            Dvcfile(repo, stage.path).dump(stage)
-            pbar_stages.update()
-
-    return link_failures
-
-
-def _find_all_targets(repo, target, recursive):
+def _find_all_targets(
+    repo: "Repo", target: str, recursive: bool = False
+) -> Iterator[str]:
     from dvc.dvcfile import is_dvc_file
 
     if os.path.isdir(target) and recursive:
-        return [
+        yield from (
             os.fspath(path)
-            for path in Tqdm(
-                repo.dvcignore.walk_files(repo.fs, target),
-                desc="Searching " + target,
-                bar_format=Tqdm.BAR_FMT_NOTOTAL,
-                unit="file",
-            )
+            for path in repo.dvcignore.walk_files(repo.fs, target)
             if not repo.is_dvc_internal(os.fspath(path))
             if not is_dvc_file(os.fspath(path))
             if not repo.scm.belongs_to_scm(os.fspath(path))
             if not repo.scm.is_tracked(os.fspath(path))
-        ]
-    return [target]
-
-
-def _create_stages(
-    repo,
-    targets,
-    fname,
-    pbar=None,
-    external=False,
-    glob=False,
-    desc=None,
-    transfer=False,
-    **kwargs,
-):
-    from dvc.stage import Stage, create_stage, restore_meta
-
-    expanded_targets = glob_targets(targets, glob=glob)
-
-    stages = []
-    for out in Tqdm(
-        expanded_targets,
-        desc="Creating DVC files",
-        disable=len(expanded_targets) < LARGE_DIR_SIZE,
-        unit="file",
-    ):
-        if kwargs.get("out"):
-            out = resolve_output(out, kwargs["out"])
-        path, wdir, out = resolve_paths(
-            repo, out, always_local=transfer and not kwargs.get("out")
         )
-        stage = create_stage(
-            Stage,
-            repo,
-            fname or path,
+    else:
+        yield target
+
+
+def create_stages(
+    repo: "Repo",
+    targets: Iterator[str],
+    fname: str = None,
+    transfer: bool = False,
+    external: bool = False,
+    **kwargs: Any,
+) -> Iterator["Stage"]:
+    for target in targets:
+        if kwargs.get("out"):
+            target = resolve_output(target, kwargs["out"])
+        path, wdir, out = resolve_paths(
+            repo, target, always_local=transfer and not kwargs.get("out")
+        )
+
+        stage = repo.stage.create(
+            single_stage=True,
+            validate=False,
+            fname=fname or path,
             wdir=wdir,
             outs=[out],
             external=external,
         )
-        restore_meta(stage)
-        if desc:
-            stage.outs[0].desc = desc
-
-        if not stage:
-            if pbar is not None:
-                pbar.total -= 1
-            continue
-
-        stages.append(stage)
-        if pbar is not None:
-            pbar.update_msg(out)
-
-    return stages
+        if kwargs.get("desc"):
+            stage.outs[0].desc = kwargs["desc"]
+        yield stage

dvc/stage/__init__.py

@@ -4,7 +4,7 @@ import string
 from collections import defaultdict
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, Optional, Set
+from typing import TYPE_CHECKING, Any, Dict, Optional, Set
 
 from funcy import cached_property, project
 
@@ -390,6 +390,19 @@ class Stage(params.StageParams):
         if purge:
             self.dvcfile.remove_stage(self)
 
+    def transfer(
+        self,
+        source: str,
+        odb: "ObjectDB" = None,
+        to_remote: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        assert len(self.outs) == 1
+        (out,) = self.outs
+        out.transfer(source, odb=odb, jobs=kwargs.get("jobs"))
+        if not to_remote:
+            out.checkout()
+
     @rwlocked(read=["deps"], write=["outs"])
     def reproduce(self, interactive=False, **kwargs):
         if not (kwargs.get("force", False) or self.changed()):

dvc/ui/__init__.py

@@ -13,6 +13,8 @@ from typing import (
 from funcy import cached_property
 
 if TYPE_CHECKING:
+    from rich.status import Status
+
     from dvc.progress import Tqdm
     from dvc.ui.table import Headers, Styles, TableData
 
@@ -181,6 +183,9 @@ class Console:
             self, data, headers, markdown=markdown, pager=pager, force=force
         )
 
+    def status(self, status: str, **kwargs: Any) -> "Status":
+        return self.error_console.status(status, **kwargs)
+
 
 ui = Console()
 

tests/func/test_add.py

@@ -1,6 +1,5 @@
 import errno
 import filecmp
-import logging
 import os
 import shutil
 import stat
@@ -131,14 +130,12 @@ class TestAddCmdDirectoryRecursive(TestDvc):
 
     def test_warn_about_large_directories(self):
         warning = (
-            "You are adding a large directory 'large-dir' recursively,"
-            " consider tracking it as a whole instead.\n"
-            "{purple}HINT:{nc} Remove the generated DVC file and then"
-            " run `{cyan}dvc add large-dir{nc}`".format(
-                purple=colorama.Fore.MAGENTA,
-                cyan=colorama.Fore.CYAN,
-                nc=colorama.Style.RESET_ALL,
-            )
+            "You are adding a large directory 'large-dir' recursively."
+            "\nConsider tracking it as a whole instead with "
+            "`{cyan}dvc add large-dir{nc}`"
+        ).format(
+            cyan=colorama.Fore.CYAN,
+            nc=colorama.Style.RESET_ALL,
         )
 
         os.mkdir("large-dir")
@@ -149,9 +146,8 @@ class TestAddCmdDirectoryRecursive(TestDvc):
             with open(path, "w") as fobj:
                 fobj.write(path)
 
-        with self._caplog.at_level(logging.WARNING, logger="dvc"):
-            assert main(["add", "--recursive", "large-dir"]) == 0
-            assert warning in self._caplog.messages
+        assert main(["add", "--recursive", "large-dir"]) == 0
+        assert warning in self._capsys.readouterr()[1]
 
 
 class TestAddDirectoryWithForwardSlash(TestDvc):
@@ -643,7 +639,7 @@ def test_failed_add_cleanup(tmp_dir, scm, dvc):
 
 
 def test_should_not_track_git_internal_files(mocker, dvc, tmp_dir):
-    stage_creator_spy = mocker.spy(dvc_module.repo.add, "_create_stages")
+    stage_creator_spy = mocker.spy(dvc_module.repo.add, "create_stages")
 
     ret = main(["add", "-R", dvc.root_dir])
     assert ret == 0
@@ -951,15 +947,15 @@ def test_add_file_in_symlink_dir(make_tmp_dir, tmp_dir, dvc, external):
         dvc.add(os.path.join("dir", "foo"))
 
 
-def test_add_with_cache_link_error(tmp_dir, dvc, mocker, caplog):
+def test_add_with_cache_link_error(tmp_dir, dvc, mocker, capsys):
     tmp_dir.gen("foo", "foo")
 
     mocker.patch(
         "dvc.checkout._do_link", side_effect=DvcException("link failed")
     )
-    with caplog.at_level(logging.WARNING, logger="dvc"):
-        dvc.add("foo")
-        assert "reconfigure cache types" in caplog.text
+    dvc.add("foo")
+    err = capsys.readouterr()[1]
+    assert "reconfigure cache types" in err
 
     assert not (tmp_dir / "foo").exists()
     assert (tmp_dir / "foo.dvc").exists()