stage: implement add/create methods in stage loader (#5397)

* stage: add/create methods in stage loader

* minor fixups

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
 
 

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()

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,

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

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("=")

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

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}'"

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",

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()
 

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",

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": []},
+    ]