From cd37ebca3753d016420aa05a96c1f23d5d739ab8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?=
 =?UTF-8?q?=E0=A4=A4=29?= <suagatchhetri@outlook.com>
Date: Tue, 7 Dec 2021 08:33:32 +0545
Subject: [PATCH] exp init: minor refactoring

---
 dvc/command/experiments.py             | 11 +++++++----
 dvc/repo/experiments/init.py           |  9 ++++-----
 tests/unit/command/test_experiments.py | 15 ++++++++-------
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py
index b097e9e12..70e310341 100644
--- a/dvc/command/experiments.py
+++ b/dvc/command/experiments.py
@@ -793,6 +793,7 @@ class CmdExperimentsRemove(CmdBase):
 
 
 class CmdExperimentsInit(CmdBase):
+    DEFAULT_NAME = "train"
     CODE = "src"
     DATA = "data"
     MODELS = "models"
@@ -851,12 +852,10 @@ class CmdExperimentsInit(CmdBase):
             force=self.args.force,
         )
 
-        name = self.args.name or "train"
-
         text = ui.rich_text.assemble(
             "\n" if self.args.interactive else "",
             "Created ",
-            (name, "bright_blue"),
+            (self.args.name, "bright_blue"),
             " stage in ",
             ("dvc.yaml", "green"),
             ".",
@@ -1452,7 +1451,11 @@ def add_parser(subparsers, parent_parser):
         help="Only use the path values explicitly provided",
     )
     experiments_init_parser.add_argument(
-        "--name", "-n", help="Name of the stage to create"
+        "--name",
+        "-n",
+        help="Name of the stage to create "
+        f"(default: {CmdExperimentsInit.DEFAULT_NAME})",
+        default=CmdExperimentsInit.DEFAULT_NAME,
     )
     experiments_init_parser.add_argument(
         "--code",
diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py
index d0aaa4786..0b4fb0920 100644
--- a/dvc/repo/experiments/init.py
+++ b/dvc/repo/experiments/init.py
@@ -29,7 +29,6 @@ from dvc.utils.serialize import dumps_yaml
 if TYPE_CHECKING:
     from dvc.repo import Repo
     from dvc.dvcfile import DVCFile
-    from dvc.stage import Stage
     from rich.tree import Tree
 
 from dvc.ui import ui
@@ -202,19 +201,17 @@ def validate_prompts(key: str, value: str) -> Union[Any, Tuple[Any, str]]:
 
 def init(
     repo: "Repo",
-    name: str = None,
+    name: str = "train",
     type: str = "default",  # pylint: disable=redefined-builtin
     defaults: Dict[str, str] = None,
     overrides: Dict[str, str] = None,
     interactive: bool = False,
     force: bool = False,
     stream: Optional[TextIO] = None,
-) -> "Stage":
+) -> PipelineStage:
     from dvc.dvcfile import make_dvcfile
 
     dvcfile = make_dvcfile(repo, "dvc.yaml")
-    name = name or "train"
-
     _check_stage_exists(dvcfile, name, force=force)
 
     defaults = defaults.copy() if defaults else {}
@@ -284,4 +281,6 @@ def init(
                 repo.scm_context.track_file(params)
     else:
         raise DvcException("Aborting ...")
+
+    assert isinstance(stage, PipelineStage)
     return stage
diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py
index 8a3681d7c..b57a6e329 100644
--- a/tests/unit/command/test_experiments.py
+++ b/tests/unit/command/test_experiments.py
@@ -613,7 +613,7 @@ def test_experiments_init(dvc, scm, mocker, capsys, extra_args):
     assert cmd.run() == 0
     m.assert_called_once_with(
         ANY(Repo),
-        name=None,
+        name="train",
         type="default",
         defaults={
             "code": "src",
@@ -653,7 +653,7 @@ def test_experiments_init_config(dvc, scm, mocker):
     assert cmd.run() == 0
     m.assert_called_once_with(
         ANY(Repo),
-        name=None,
+        name="train",
         type="default",
         defaults={
             "code": "new_src",
@@ -678,7 +678,7 @@ def test_experiments_init_explicit(dvc, mocker):
     assert cmd.run() == 0
     m.assert_called_once_with(
         ANY(Repo),
-        name=None,
+        name="train",
         type="default",
         defaults={},
         overrides={"cmd": "cmd"},
@@ -710,9 +710,9 @@ def test_experiments_init_cmd_not_required_for_interactive_mode(dvc, mocker):
 @pytest.mark.parametrize(
     "extra_args, expected_kw",
     [
-        (["--type", "default"], {"type": "default"}),
-        (["--type", "dl"], {"type": "dl"}),
-        (["--force"], {"force": True}),
+        (["--type", "default"], {"type": "default", "name": "train"}),
+        (["--type", "dl"], {"type": "dl", "name": "train"}),
+        (["--force"], {"force": True, "name": "train"}),
         (
             ["--name", "name", "--type", "dl"],
             {"name": "name", "type": "dl"},
@@ -735,6 +735,7 @@ def test_experiments_init_cmd_not_required_for_interactive_mode(dvc, mocker):
                 "live",
             ],
             {
+                "name": "train",
                 "overrides": {
                     "plots": "p",
                     "models": "m",
@@ -744,7 +745,7 @@ def test_experiments_init_cmd_not_required_for_interactive_mode(dvc, mocker):
                     "data": "d",
                     "live": "live",
                     "cmd": "cmd",
-                }
+                },
             },
         ),
     ],