diff --git a/dvc/repo/add.py b/dvc/repo/add.py index b7f09658b..3e48451c2 100644 --- a/dvc/repo/add.py +++ b/dvc/repo/add.py @@ -14,6 +14,7 @@ from ..exceptions import ( 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: @@ -23,34 +24,29 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -@locked -@scm_context -def add( # noqa: C901 - repo, - targets: "TargetType", - recursive=False, - no_commit=False, - fname=None, - to_remote=False, - **kwargs, -): - from dvc.utils.collections import ensure_list - - if recursive and fname: +def check_recursive_and_fname(args): + if args.recursive and args.fname: raise RecursiveAddingWhileUsingFilename() - targets = ensure_list(targets) - to_cache = kwargs.get("out") and not to_remote +def transform_targets(args): + args.targets = ensure_list(args.targets) + + +def check_arg_combinations(args): + kwargs = args.kwargs invalid_opt = None + to_remote = args.to_remote + to_cache = kwargs.get("out") and not to_remote + if to_remote or to_cache: message = "{option} can't be used with " message += "--to-remote" if to_remote else "-o" - if len(targets) != 1: + if len(args.targets) != 1: invalid_opt = "multiple targets" - elif no_commit: + elif args.no_commit: invalid_opt = "--no-commit option" - elif recursive: + elif args.recursive: invalid_opt = "--recursive option" elif kwargs.get("external"): invalid_opt = "--external option" @@ -64,9 +60,30 @@ def add( # noqa: C901 if invalid_opt is not None: raise InvalidArgumentError(message.format(option=invalid_opt)) + +VALIDATORS = ( + check_recursive_and_fname, + transform_targets, + check_arg_combinations, +) + + +@locked +@scm_context +@validate(*VALIDATORS) +def add( # noqa: C901 + repo, + targets: "TargetType", + recursive=False, + no_commit=False, + fname=None, + to_remote=False, + **kwargs, +): 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 diff --git a/dvc/utils/collections.py b/dvc/utils/collections.py index 84cee0de5..ad50224ec 100644 --- a/dvc/utils/collections.py +++ b/dvc/utils/collections.py @@ -1,6 +1,8 @@ +import inspect import os from collections.abc import Mapping -from typing import Dict, Iterable, List, TypeVar, Union +from functools import wraps +from typing import Callable, Dict, Iterable, List, TypeVar, Union from pygtrie import StringTrie as _StringTrie @@ -76,3 +78,59 @@ def chunk_dict(d: Dict[_KT, _VT], size: int = 1) -> List[Dict[_KT, _VT]]: from funcy import chunks return [{key: d[key] for key in chunk} for chunk in chunks(size, d)] + + +class _NamespacedDict(dict): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.__dict__ = self + + +def validate(*validators: Callable, post: bool = False): + """ + Validate and transform arguments and results from function calls. + + The validators functions are passed a dictionary of arguments, which + supports dot notation access too. + + The key is derived from the function signature, and hence is the name of + the argument, whereas the value is the one passed to the function + (if it is not passed, default value from keyword arguments are provided). + + >>> def validator(args): + ... assert args["l"] >= 0 and args.b >= 0 and args.h >= 0 + + >>> @validate(validator) + ... def cuboid_area(l, b, h=1): + ... return 2*(l*b + l*h + b*h) + + >>> cuboid_area(5, 20) + 250 + >>> cuboid_area(-1, -2) + Traceback (most recent call last): + ... + AssertionError + """ + + def wrapped(func: Callable): + sig = inspect.signature(func) + + @wraps(func) + def inner(*args, **kwargs): + ba = sig.bind(*args, **kwargs) + ba.apply_defaults() + ba.arguments = _NamespacedDict(ba.arguments) + + if not post: + for validator in validators: + validator(ba.arguments) + + result = func(*ba.args, **ba.kwargs) + if post: + for validator in validators: + result = validator(result) + return result + + return inner + + return wrapped diff --git a/tests/unit/utils/test_collections.py b/tests/unit/utils/test_collections.py index 2c20b9a5f..9de4b1917 100644 --- a/tests/unit/utils/test_collections.py +++ b/tests/unit/utils/test_collections.py @@ -1,5 +1,7 @@ # pylint: disable=unidiomatic-typecheck -from dvc.utils.collections import apply_diff, chunk_dict +from mock import create_autospec + +from dvc.utils.collections import apply_diff, chunk_dict, validate class MyDict(dict): @@ -58,3 +60,78 @@ def test_chunk_dict(): assert chunk_dict(d, 2) == [{"a": 1, "b": 2}, {"c": 3}] assert chunk_dict(d, 3) == [d] assert chunk_dict(d, 4) == [d] + + +# pylint: disable=unused-argument + + +def _test_func(x, y, *args, j=3, k=5, **kwargs): + pass + + +def test_pre_validate_decorator_required_args(mocker): + mock = mocker.MagicMock() + + func_mock = create_autospec(_test_func) + func = validate(mock)(func_mock) + + func("x", "y") + + func_mock.assert_called_once_with("x", "y", j=3, k=5) + mock.assert_called_once() + + (args,) = mock.call_args[0] + assert args.x == "x" + assert args.y == "y" + assert args.args == () + assert args.j == 3 + assert args.k == 5 + assert args.kwargs == {} + + +def test_pre_validate_decorator_kwargs_args(mocker): + mock = mocker.MagicMock() + func_mock = create_autospec(_test_func) + func = validate(mock)(func_mock) + + func("x", "y", "part", "of", "args", j=1, k=10, m=100, n=1000) + + func_mock.assert_called_once_with( + "x", "y", "part", "of", "args", j=1, k=10, m=100, n=1000 + ) + mock.assert_called_once() + (args,) = mock.call_args[0] + assert args.x == "x" + assert args.y == "y" + assert args.args == ("part", "of", "args") + assert args.j == 1 + assert args.k == 10 + assert args.kwargs == {"m": 100, "n": 1000} + + +def test_pre_validate_update_args(): + def test_validator(args): + args.w += 50 + del args.x + args.y = 100 + + def test_func(w=1, x=5, y=10, z=15): + pass + + mock = create_autospec(test_func) + func = validate(test_validator)(mock) + + func(100, 100) + mock.assert_called_once_with(w=150, y=100, z=15) + + +def test_post_validate_decorator(mocker): + def none_filter(result): + return list(filter(None, result)) + + test_func = mocker.MagicMock(return_value=[1, None, 2]) + func = validate(none_filter, post=True)(test_func) + + result = func() + test_func.assert_called_once() + assert result == [1, 2]