exp init: refactor and simplify interactive mode

dvc/repo/experiments/init.py

@@ -12,19 +12,14 @@
     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)

tests/func/experiments/test_init.py

@@ -5,7 +5,6 @@
 
 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