diff --git a/dvc/repo/add.py b/dvc/repo/add.py index 24579a36c..93e044a96 100644 --- a/dvc/repo/add.py +++ b/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 diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 8559b615a..c45d7df6c 100644 --- a/dvc/stage/__init__.py +++ b/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()): diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py index 71d42e798..441b5b929 100644 --- a/dvc/ui/__init__.py +++ b/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() diff --git a/tests/func/test_add.py b/tests/func/test_add.py index 071a2f753..9273951e2 100644 --- a/tests/func/test_add.py +++ b/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()