exp init: iterate on ui, tree and show to-be-generated yaml content (#6750)

* iterate on ui, has tree and shows to-be-generated yaml content

* order fs tree structure alphabetically

* always show the heading

* hide dvclive on tree if not on live mode

* Change message at the end

* s/skip/omit

* add link on the header

* remove newline on non-interactive mode

dvc/command/experiments.py

@@ -806,6 +806,10 @@ class CmdExperimentsInit(CmdBase):
         "plots": PLOTS,
         "live": DVCLIVE,
     }
+    EXP_LINK = (
+        "https://dvc.org/doc"
+        "/user-guide/experiment-management/running-experiments"
+    )
 
     def run(self):
         from dvc.command.stage import parse_cmd
@@ -843,10 +847,34 @@ class CmdExperimentsInit(CmdBase):
             interactive=self.args.interactive,
             force=self.args.force,
         )
+
+        name = self.args.name or self.args.type
+
+        text = ui.rich_text.assemble(
+            "\n" if self.args.interactive else "",
+            "Created ",
+            (name, "bright_blue"),
+            " stage in ",
+            ("dvc.yaml", "green"),
+            ".",
+        )
+        if not self.args.run:
+            text.append_text(
+                ui.rich_text.assemble(
+                    " To run, use ",
+                    ('"dvc exp run"', "green"),
+                    ".\nSee ",
+                    (self.EXP_LINK, "repr.url"),
+                    ".",
+                )
+            )
+
+        ui.write(text, styled=True)
         if self.args.run:
             return self.repo.experiments.run(
                 targets=[initialized_stage.addressing]
             )
+
         return 0
 
 

dvc/repo/experiments/init.py

@@ -1,10 +1,18 @@
+import logging
 import os
-from typing import TYPE_CHECKING, Dict, Iterable
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Dict, Iterable, cast
 
 from funcy import compact, lremove, post_processing
 from rich.prompt import Prompt
+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.serialize import dumps_yaml
 
 if TYPE_CHECKING:
     from dvc.repo import Repo
@@ -71,7 +79,7 @@ class SkippablePrompt(RequiredPrompt):
             prompt.append(_default)
             prompt.append(", ")
 
-        prompt.append(f"{self.skip_value} to skip", style="italic")
+        prompt.append(f"{self.skip_value} to omit", style="italic")
         prompt.append("]")
         prompt.append(self.prompt_suffix)
         return prompt
@@ -92,24 +100,67 @@ def _prompts(keys: Iterable[str], defaults: Dict[str, OptStr]):
         yield key, value
 
 
+@contextmanager
+def _disable_logging(highest_level=logging.CRITICAL):
+    previous_level = logging.root.manager.disable
+
+    logging.disable(highest_level)
+
+    try:
+        yield
+    finally:
+        logging.disable(previous_level)
+
+
+PIPELINE_FILE_LINK = (
+    "https://dvc.org/doc/user-guide/project-structure/pipelines-files"
+)
+
+
 def init_interactive(
+    name: str,
     defaults: Dict[str, str],
-    provided: Iterable[str],
-    show_heading: bool = False,
+    provided: Dict[str, str],
+    show_tree: bool = False,
     live: bool = False,
 ) -> Dict[str, str]:
-    primary = lremove(provided, ["cmd", "code", "data", "models", "params"])
-    secondary = lremove(provided, ["live"] if live else ["metrics", "plots"])
+    primary = lremove(
+        provided.keys(), ["cmd", "code", "data", "models", "params"]
+    )
+    secondary = lremove(
+        provided.keys(), ["live"] if live else ["metrics", "plots"]
+    )
 
     if not (primary or secondary):
         return {}
 
-    message = (
-        "This command will guide you to set up your first stage in "
-        "[green]dvc.yaml[/green].\n"
+    message = ui.rich_text.assemble(
+        "This command will guide you to set up a ",
+        (name, "bright_blue"),
+        " stage in ",
+        ("dvc.yaml", "green"),
+        ".",
     )
-    if show_heading:
-        ui.error_write(message, styled=True)
+    doc_link = ui.rich_text.assemble(
+        "See ", (PIPELINE_FILE_LINK, "repr.url"), "."
+    )
+    ui.error_write(message, doc_link, "", sep="\n", styled=True)
+
+    if show_tree:
+        from rich.tree import Tree
+
+        tree = Tree(
+            "DVC assumes the following workspace structure:",
+            highlight=True,
+        )
+        workspace = {**defaults, **provided}
+        workspace.pop("cmd", None)
+        if not live and "live" not in provided:
+            workspace.pop("live", None)
+        for value in sorted(workspace.values()):
+            tree.add(f"[green]{value}[/green]")
+        ui.error_write(tree, styled=True)
+        ui.error_write()
 
     return compact(
         {
@@ -153,10 +204,11 @@ def init(
     with_live = type == "live"
     if interactive:
         defaults = init_interactive(
-            defaults=defaults or {},
-            show_heading=not dvcfile.exists(),
+            name,
+            defaults=defaults,
             live=with_live,
-            provided=overrides.keys(),
+            provided=overrides,
+            show_tree=True,
         )
     else:
         if with_live:
@@ -182,7 +234,7 @@ def init(
 
     checkpoint_out = bool(context.get("live"))
     models = context.get("models")
-    return repo.stage.add(
+    stage = repo.stage.create(
         name=name,
         cmd=context["cmd"],
         deps=compact([context.get("code"), context.get("data")]),
@@ -193,3 +245,19 @@ def init(
         force=force,
         **{"checkpoints" if checkpoint_out else "outs": compact([models])},
     )
+
+    if interactive:
+        ui.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)
+
+    if not interactive or ui.confirm(
+        "Do you want to add the above contents to dvc.yaml?"
+    ):
+        with _disable_logging():
+            stage.dump(update_lock=False)
+        stage.ignore_outs()
+    else:
+        raise DvcException("Aborting ...")
+    return stage