diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py index 1a731612b..85f2e6b01 100644 --- a/dvc/repo/experiments/init.py +++ b/dvc/repo/experiments/init.py @@ -12,19 +12,14 @@ from typing import ( TextIO, Tuple, Union, - cast, ) from funcy import compact, lremove, lsplit -from rich.rule import Rule -from rich.syntax import Syntax from dvc.exceptions import DvcException from dvc.stage import PipelineStage -from dvc.stage.serialize import to_pipeline_file from dvc.types import OptStr from dvc.utils import humanize -from dvc.utils.serialize import dumps_yaml if TYPE_CHECKING: from dvc.repo import Repo @@ -89,15 +84,13 @@ def build_workspace_tree(workspace: Dict[str, str], label: str) -> "Tree": return tree -def display_workspace_tree( - workspace: Dict[str, str], label: str, stderr: bool = False -) -> None: +def display_workspace_tree(workspace: Dict[str, str], label: str) -> None: d = workspace.copy() d.pop("cmd", None) if d: - ui.write(build_workspace_tree(d, label), styled=True, stderr=stderr) - ui.write(styled=True, stderr=stderr) + ui.write(build_workspace_tree(d, label), styled=True) + ui.write(styled=True) PIPELINE_FILE_LINK = "https://s.dvc.org/g/pipeline-files" @@ -112,55 +105,40 @@ def init_interactive( stream: Optional[TextIO] = None, ) -> Dict[str, str]: command = provided.pop("cmd", None) - primary = lremove(provided.keys(), ["code", "data", "models", "params"]) - secondary = lremove( - provided.keys(), ["live"] if live else ["metrics", "plots"] + prompts = lremove(provided.keys(), ["code", "data", "models", "params"]) + prompts.extend( + lremove(provided.keys(), ["live"] if live else ["metrics", "plots"]) ) - prompts = primary + secondary - - workspace = {**defaults, **provided} - if not live and "live" not in provided: - workspace.pop("live", None) - for key in ("plots", "metrics"): - if live and key not in provided: - workspace.pop(key, None) ret: Dict[str, str] = {} + if prompts or command: + ui.error_write( + f"This command will guide you to set up a [bright_blue]{name}[/]", + "stage in [green]dvc.yaml[/].", + f"\nSee [repr.url]{PIPELINE_FILE_LINK}[/].\n", + styled=True, + ) + if command: ret["cmd"] = command - - if not prompts and command: - return ret - - ui.error_write( - f"This command will guide you to set up a [bright_blue]{name}[/]", - "stage in [green]dvc.yaml[/].", - f"\nSee [repr.url]{PIPELINE_FILE_LINK}[/].\n", - styled=True, - ) - - if not command: + else: ret.update( compact(_prompts(["cmd"], allow_omission=False, stream=stream)) ) if prompts: ui.error_write(styled=True) - if not prompts: - return ret - - ui.error_write( - "Enter the paths for dependencies and outputs of the command.", - styled=True, - ) - - tree_label = "Using experiment project structure:" - display_workspace_tree(workspace, tree_label, stderr=True) - ret.update( - compact( - _prompts(prompts, defaults, validator=validator, stream=stream) + if prompts: + ui.error_write( + "Enter the paths for dependencies and outputs of the command.", + styled=True, ) - ) + ret.update( + compact( + _prompts(prompts, defaults, validator=validator, stream=stream) + ) + ) + ui.error_write(styled=True) return ret @@ -311,31 +289,10 @@ def init( **{"checkpoints" if checkpoint_out else "outs": compact([models])}, ) - if interactive: - ui.error_write(Rule(style="green"), styled=True) - _yaml = dumps_yaml(to_pipeline_file(cast(PipelineStage, stage))) - syn = Syntax(_yaml, "yaml", theme="ansi_dark") - ui.error_write(syn, styled=True) - - from dvc.ui.prompt import Confirm - - if interactive and not Confirm.ask( - "Do you want to add the above contents to dvc.yaml?", - console=ui.error_console, - default=True, - stream=stream, - ): - raise DvcException("Aborting ...") - - if interactive: - ui.error_write() # add a newline after the prompts - with _disable_logging(), repo.scm_context(autostage=True, quiet=True): stage.dump(update_lock=False) stage.ignore_outs() - if not interactive: - label = "Using experiment project structure:" - display_workspace_tree(context, label) + display_workspace_tree(context, "Using experiment project structure:") init_deps(stage) if params: repo.scm_context.track_file(params) diff --git a/tests/func/experiments/test_init.py b/tests/func/experiments/test_init.py index 2a366417e..fc0bc61d8 100644 --- a/tests/func/experiments/test_init.py +++ b/tests/func/experiments/test_init.py @@ -5,7 +5,6 @@ import pytest from dvc.cli import main from dvc.commands.experiments.init import CmdExperimentsInit -from dvc.exceptions import DvcException from dvc.repo.experiments.init import init from dvc.stage.exceptions import DuplicateStageName @@ -95,21 +94,6 @@ def test_init_with_no_defaults_non_interactive(tmp_dir, scm, dvc): assert scm.is_tracked("dvc.yaml") -def test_abort_confirmation(tmp_dir, dvc): - (tmp_dir / "param").dump({"foo": 1}) - inp = io.StringIO("./script\nscript\ndata\nmodel\nparam\nmetric\nplt\nn") - with pytest.raises(DvcException) as exc: - init( - dvc, - interactive=True, - defaults=CmdExperimentsInit.DEFAULTS, - stream=inp, - ) - assert str(exc.value) == "Aborting ..." - assert not (tmp_dir / "dvc.yaml").exists() - assert not (tmp_dir / "dvc.lock").exists() - - @pytest.mark.parametrize( "extra_overrides, inp", [ @@ -213,7 +197,7 @@ def test_init_interactive_params_validation(tmp_dir, dvc, capsys): "Please retry with an existing parameters file.\n" "Path to a parameters file [params.yaml, n to omit]:" ) in err - assert out == "Created script.py.\n" + assert "Created script.py." in out def test_init_with_no_defaults_interactive(tmp_dir, dvc): @@ -296,8 +280,7 @@ def test_init_default(tmp_dir, scm, dvc, interactive, overrides, inp, capsys): if interactive: assert "'script.py' does not exist, the file will be created." in err assert "'data' does not exist, the directory will be created." in err - else: - assert "Using experiment project structure: " in out + assert "Using experiment project structure: " in out assert "Created script.py and data" in out @@ -382,8 +365,7 @@ def test_init_interactive_live( if interactive: assert "'script.py' does not exist, the file will be created." in err assert "'data' does not exist, the directory will be created." in err - else: - assert "Using experiment project structure: " in out + assert "Using experiment project structure: " in out assert "Created script.py and data" in out