diff --git a/dvc/command/stage.py b/dvc/command/stage.py index 50bacce51..a86fa1a8d 100644 --- a/dvc/command/stage.py +++ b/dvc/command/stage.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Dict, Iterable, List from dvc.command import completion from dvc.command.base import CmdBase, append_doc_link, fix_subparsers +from dvc.utils.cli_parse import parse_params if TYPE_CHECKING: from dvc.output.base import BaseOutput @@ -127,15 +128,14 @@ def parse_cmd(commands: List[str]) -> str: class CmdStageAdd(CmdBase): def run(self): - repo = self.repo kwargs = vars(self.args) - kwargs["cmd"] = parse_cmd(kwargs.pop("cmd")) - - stage = repo.stage.create_from_cli(validate=True, **kwargs) - - with repo.scm.track_file_changes(config=repo.config): - stage.ignore_outs() - stage.dump(update_lock=False) + kwargs.update( + { + "cmd": parse_cmd(kwargs.pop("cmd")), + "params": parse_params(self.args.params), + } + ) + self.repo.stage.add(**kwargs) return 0 diff --git a/dvc/repo/run.py b/dvc/repo/run.py index 06fb1bb5d..c44bf2558 100644 --- a/dvc/repo/run.py +++ b/dvc/repo/run.py @@ -1,5 +1,7 @@ from typing import TYPE_CHECKING, Union +from dvc.utils.cli_parse import parse_params + from . import locked from .scm_context import scm_context @@ -18,12 +20,12 @@ def run( run_cache: bool = True, force: bool = True, **kwargs -) -> Union["Stage", "PipelineStage", None]: - from dvc.stage.utils import validate_state +) -> Union["Stage", "PipelineStage"]: - stage = self.stage.create_from_cli(**kwargs) - - validate_state(self, stage, force=force) + kwargs.update( + {"force": force, "params": parse_params(kwargs.get("params", []))} + ) + stage = self.stage.create(**kwargs) if no_exec: stage.ignore_outs() diff --git a/dvc/repo/stage.py b/dvc/repo/stage.py index 0bb3b45c3..cacae58f6 100644 --- a/dvc/repo/stage.py +++ b/dvc/repo/stage.py @@ -3,7 +3,7 @@ import logging import os import typing from contextlib import suppress -from functools import partial +from functools import partial, wraps from typing import ( Callable, Iterable, @@ -21,6 +21,7 @@ from dvc.exceptions import ( OutputNotFoundError, ) from dvc.path_info import PathInfo +from dvc.repo import lock_repo from dvc.utils import parse_target, relpath logger = logging.getLogger(__name__) @@ -99,24 +100,94 @@ def _collect_specific_target( return [], file, name +def locked(f): + @wraps(f) + def wrapper(loader: "StageLoad", *args, **kwargs): + with lock_repo(loader.repo): + return f(loader, *args, **kwargs) + + return wrapper + + class StageLoad: def __init__(self, repo: "Repo") -> None: - self.repo = repo + self.repo: "Repo" = repo - def create_from_cli( - self, validate: bool = False, **kwargs + @locked + def add( + self, + single_stage: bool = False, + fname: str = None, + validate: bool = True, + force: bool = False, + update_lock: bool = False, + **stage_data, + ): + stage = self.create( + single_stage=single_stage, + fname=fname, + validate=validate, + force=force, + **stage_data, + ) + scm = self.repo.scm + with scm.track_file_changes(config=self.repo.config): + stage.dump(update_lock=update_lock) + stage.ignore_outs() + + return stage + + def create( + self, + single_stage: bool = False, + validate: bool = True, + fname: str = None, + force: bool = False, + **stage_data, ) -> Union["Stage", "PipelineStage"]: - """Creates a stage from CLI args passed as kwargs. + """Creates a stage. Args: + single_stage: if true, the .dvc file based stage is created, + fname is required in that case + fname: name of the file to use, not used for dvc.yaml files validate: if true, the new created stage is checked against the stages in the repo. Eg: graph correctness, potential overwrites in dvc.yaml file (unless `force=True`). + force: ignores overwrites in dvc.yaml file + stage_data: Stage data to create from + (see create_stage and loads_from for more information) """ - from dvc.stage.utils import create_stage_from_cli + from dvc.stage import PipelineStage, Stage, create_stage, restore_meta + from dvc.stage.exceptions import InvalidStageName + from dvc.stage.utils import ( + is_valid_name, + prepare_file_path, + validate_kwargs, + validate_state, + ) - kwargs.update(validate=validate) - return create_stage_from_cli(self.repo, **kwargs) + stage_data = validate_kwargs( + single_stage=single_stage, fname=fname, **stage_data + ) + if single_stage: + stage_cls = Stage + path = fname or prepare_file_path(stage_data) + else: + path = PIPELINE_FILE + stage_cls = PipelineStage + stage_name = stage_data["name"] + if not (stage_name and is_valid_name(stage_name)): + raise InvalidStageName + + stage = create_stage( + stage_cls, repo=self.repo, path=path, **stage_data + ) + if validate: + validate_state(self.repo, stage, force=force) + + restore_meta(stage) + return stage def from_target( self, target: str, accept_group: bool = False, glob: bool = False, diff --git a/dvc/stage/utils.py b/dvc/stage/utils.py index f48ca2119..2c327b1a4 100644 --- a/dvc/stage/utils.py +++ b/dvc/stage/utils.py @@ -1,17 +1,14 @@ import os import pathlib from contextlib import suppress -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, Union from funcy import concat, first, lsplit, rpartial, without from dvc.exceptions import InvalidArgumentError -from dvc.utils.cli_parse import parse_params -from dvc.utils.collections import chunk_dict from ..hash_info import HashInfo from .exceptions import ( - InvalidStageName, MissingDataSource, StageExternalOutputsError, StagePathNotDirectoryError, @@ -252,7 +249,7 @@ def is_valid_name(name: str): return not INVALID_STAGENAME_CHARS & set(name) -def _get_file_path(kwargs): +def prepare_file_path(kwargs): """Determine file path from the first output name. Used in creating .dvc files. @@ -330,19 +327,8 @@ def validate_state( repo.check_modified_graph(new_stages=[new], old_stages=stages) -def create_stage_from_cli( - repo: "Repo", - single_stage: bool = False, - fname: str = None, - validate: bool = False, - force: bool = False, - **kwargs: Any, -) -> Union["Stage", "PipelineStage"]: - - from dvc.dvcfile import PIPELINE_FILE - - from . import PipelineStage, Stage, create_stage, restore_meta - +def validate_kwargs(single_stage: bool = False, fname: str = None, **kwargs): + """Prepare, validate and process kwargs passed from cli""" cmd = kwargs.get("cmd") if not cmd: raise InvalidArgumentError("command is not specified") @@ -362,32 +348,16 @@ def create_stage_from_cli( if single_stage: kwargs.pop("name", None) - stage_cls = Stage - path = fname or _get_file_path(kwargs) - else: - path = PIPELINE_FILE - stage_cls = PipelineStage - if not (stage_name and is_valid_name(stage_name)): - raise InvalidStageName - kwargs["cmd"] = cmd[0] if isinstance(cmd, list) and len(cmd) == 1 else cmd - kwargs["live_summary"] = not kwargs.pop("live_no_summary", False) - kwargs["live_html"] = not kwargs.pop("live_no_html", False) - - live = kwargs.get("live", False) - live_no_cache = kwargs.get("live_no_cache", False) - if live and live_no_cache: + if kwargs.get("live") and kwargs.get("live_no_cache"): raise InvalidArgumentError( "cannot specify both `--live` and `--live-no-cache`" ) - params = chunk_dict(parse_params(kwargs.pop("params", []))) - stage = create_stage( - stage_cls, repo=repo, path=path, params=params, **kwargs + kwargs.update( + { + "live_summary": not kwargs.pop("live_no_summary", False), + "live_html": not kwargs.pop("live_no_html", False), + } ) - - if validate: - validate_state(repo, stage, force=force) - - restore_meta(stage) - return stage + return kwargs diff --git a/dvc/utils/cli_parse.py b/dvc/utils/cli_parse.py index 45b0e656f..3bbd597f9 100644 --- a/dvc/utils/cli_parse.py +++ b/dvc/utils/cli_parse.py @@ -2,7 +2,7 @@ from collections import defaultdict from typing import Any, Dict, Iterable, List -def parse_params(path_params: Iterable[str]) -> Dict[str, List[str]]: +def parse_params(path_params: Iterable[str]) -> List[Dict[str, List[str]]]: """Normalizes the shape of params from the CLI to dict.""" from dvc.dependency.param import ParamsDependency @@ -14,7 +14,7 @@ def parse_params(path_params: Iterable[str]) -> Dict[str, List[str]]: if not path: path = ParamsDependency.DEFAULT_PARAMS_FILE ret[path].extend(params) - return ret + return [{path: params} for path, params in ret.items()] def loads_params(path_params: Iterable[str],) -> Dict[str, Dict[str, Any]]: @@ -29,7 +29,11 @@ def loads_params(path_params: Iterable[str],) -> Dict[str, Dict[str, Any]]: normalized_params = parse_params(path_params) ret: Dict[str, Dict[str, Any]] = defaultdict(dict) - for path, param_keys in normalized_params.items(): + for part in normalized_params: + assert part + (item,) = part.items() + path, param_keys = item + for param_str in param_keys: try: key, _, value = param_str.partition("=") diff --git a/tests/conftest.py b/tests/conftest.py index f165646f6..72797fe58 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -138,18 +138,3 @@ def pytest_configure(config): enabled_remotes.discard(remote_name) if enabled: enabled_remotes.add(remote_name) - - -@pytest.fixture -def dummy_stage(tmp_dir, dvc): - def make(path="dvc.yaml", name="dummy_stage", **kwargs): - from dvc.stage import PipelineStage, create_stage - - cmd = kwargs.get("cmd", "command") - stage = create_stage( - PipelineStage, dvc, path, name=name, cmd=cmd, **kwargs - ) - stage.dump() - return stage - - return make diff --git a/tests/unit/command/test_stage.py b/tests/unit/command/test_stage.py index 6aad40e9e..d7e21aef8 100644 --- a/tests/unit/command/test_stage.py +++ b/tests/unit/command/test_stage.py @@ -61,7 +61,7 @@ def test_stage_add(mocker, dvc, command, parsed_command): assert cli_args.func == CmdStageAdd cmd = cli_args.func(cli_args) - m = mocker.patch.object(cmd.repo.stage, "create_from_cli") + m = mocker.patch.object(cmd.repo.stage, "add") assert cmd.run() == 0 expected = { @@ -69,7 +69,10 @@ def test_stage_add(mocker, dvc, command, parsed_command): "deps": ["deps"], "outs": ["outs"], "outs_no_cache": ["outs-no-cache"], - "params": ["file:param1,param2", "param3"], + "params": [ + {"file": ["param1", "param2"]}, + {"params.yaml": ["param3"]}, + ], "metrics": ["metrics"], "metrics_no_cache": ["metrics-no-cache"], "plots": ["plots"], @@ -85,6 +88,7 @@ def test_stage_add(mocker, dvc, command, parsed_command): "external": True, "desc": "description", "cmd": parsed_command, + "force": True, } args, kwargs = m.call_args assert not args @@ -92,4 +96,4 @@ def test_stage_add(mocker, dvc, command, parsed_command): for key, val in expected.items(): # expected values should be a subset of what's in the call args list assert key in kwargs - assert kwargs[key] == val + assert kwargs[key] == val, f"in key '{key}'" diff --git a/tests/unit/test_metrics.py b/tests/unit/test_metrics.py index eb9f9538b..4063e26d5 100644 --- a/tests/unit/test_metrics.py +++ b/tests/unit/test_metrics.py @@ -2,7 +2,7 @@ import json import os -def test_metrics_order(tmp_dir, dvc, dummy_stage): +def test_metrics_order(tmp_dir, dvc): tmp_dir.gen( { "p.json": json.dumps({"p1": 1}), @@ -14,11 +14,17 @@ def test_metrics_order(tmp_dir, dvc, dummy_stage): } ) - dummy_stage(metrics=["p.json", str(tmp_dir / "sub" / "p4.json")]) - dummy_stage( - path=str(tmp_dir / "sub" / "dvc.yaml"), - metrics=["p1.json", str(tmp_dir / "sub" / "p3.json")], + dvc.stage.add( + metrics=["p.json", str(tmp_dir / "sub" / "p4.json")], + cmd="cmd1", + name="stage1", ) + with (tmp_dir / "sub").chdir(): + dvc.stage.add( + metrics=[str(tmp_dir / "p1.json"), str("p3.json")], + cmd="cmd2", + name="stage2", + ) assert list(dvc.metrics.show()[""]) == [ "p.json", diff --git a/tests/unit/test_params.py b/tests/unit/test_params.py index f184e1142..44b38a4f8 100644 --- a/tests/unit/test_params.py +++ b/tests/unit/test_params.py @@ -3,7 +3,7 @@ import os from dvc.utils.serialize import dumps_yaml -def test_params_order(tmp_dir, dvc, dummy_stage): +def test_params_order(tmp_dir, dvc): tmp_dir.gen( { "params.yaml": dumps_yaml({"p1": 1}), @@ -13,19 +13,25 @@ def test_params_order(tmp_dir, dvc, dummy_stage): ) p2_path = os.path.join("sub", "p2.yaml") - sub_stage = os.path.join("sub", "dvc.yaml") - - dummy_stage(params=[{p2_path: ["p3"]}, {"p1.yaml": ["p2"]}]) - dummy_stage(path=sub_stage, params=["p1"]) + dvc.stage.add( + params=[{p2_path: ["p3"]}, {"p1.yaml": ["p2"]}], + cmd="cmd1", + name="stage1", + ) + with (tmp_dir / "sub").chdir(): + dvc.stage.add(params=["p1"], cmd="cmd2", name="stage2") # params are sorted during dumping, therefore p1 is first assert list(dvc.params.show()[""]) == ["p1.yaml", p2_path, "params.yaml"] -def test_repro_unicode(tmp_dir, dvc, dummy_stage): +def test_repro_unicode(tmp_dir, dvc): tmp_dir.gen({"settings.json": '{"Ω_value": 1}'}) - dummy_stage(params=[{"settings.json": ["Ω_value"]}]) - (stage,) = dvc.reproduce(dry=True) + stage = dvc.stage.add( + params=[{"settings.json": ["Ω_value"]}], cmd="cmd", name="stage1" + ) + assert dvc.reproduce(dry=True) == [stage] + stage.cmd = "foo" stage.dump() diff --git a/tests/unit/test_plots.py b/tests/unit/test_plots.py index c3387a5d0..b0e8f2049 100644 --- a/tests/unit/test_plots.py +++ b/tests/unit/test_plots.py @@ -2,7 +2,7 @@ import json import os -def test_plots_order(tmp_dir, dvc, dummy_stage): +def test_plots_order(tmp_dir, dvc): tmp_dir.gen( { "p.json": json.dumps([{"p1": 1}, {"p1": 2}]), @@ -14,11 +14,17 @@ def test_plots_order(tmp_dir, dvc, dummy_stage): } ) - dummy_stage(plots=["p.json", str(tmp_dir / "sub" / "p4.json")]) - dummy_stage( - path=str(tmp_dir / "sub" / "dvc.yaml"), - plots=["p1.json", str(tmp_dir / "sub" / "p3.json")], + dvc.stage.add( + plots=["p.json", str(tmp_dir / "sub" / "p4.json")], + cmd="cmd1", + name="stage1", ) + with (tmp_dir / "sub").chdir(): + dvc.stage.add( + plots=[str(tmp_dir / "p1.json"), "p3.json"], + cmd="cmd2", + name="stage2", + ) assert list(dvc.plots.show()) == [ "p.json", diff --git a/tests/unit/utils/test_cli_parse.py b/tests/unit/utils/test_cli_parse.py index b08ec9cfe..267047dc9 100644 --- a/tests/unit/utils/test_cli_parse.py +++ b/tests/unit/utils/test_cli_parse.py @@ -12,9 +12,9 @@ def test_parse_params(): "param3,", "file3:", ] - ) == { - "params.yaml": ["param1", "param2", "param10", "param3"], - "file1": ["param1", "param2", "param3", "param4"], - "file2": ["param2"], - "file3": [], - } + ) == [ + {"params.yaml": ["param1", "param2", "param10", "param3"]}, + {"file1": ["param1", "param2", "param3", "param4"]}, + {"file2": ["param2"]}, + {"file3": []}, + ]