exp init: create workspace structure by default, show workspace tree

dvc/repo/experiments/init.py

@@ -8,6 +8,7 @@ from typing import (
     Callable,
     Dict,
     Iterable,
+    List,
     Optional,
     TextIO,
     Tuple,
@@ -29,6 +30,7 @@ if TYPE_CHECKING:
     from dvc.repo import Repo
     from dvc.dvcfile import DVCFile
     from rich.tree import Tree
+    from dvc.dependency import Dependency
 
 from dvc.ui import ui
 
@@ -79,18 +81,26 @@ def _disable_logging(highest_level=logging.CRITICAL):
         logging.disable(previous_level)
 
 
-def build_workspace_tree(workspace: Dict[str, str]) -> "Tree":
+def build_workspace_tree(workspace: Dict[str, str], label: str) -> "Tree":
     from rich.tree import Tree
 
-    tree = Tree(
-        "DVC assumes the following workspace structure:",
-        highlight=True,
-    )
+    tree = Tree(label, highlight=True)
     for value in sorted(workspace.values()):
         tree.add(f"[green]{value}[/green]")
     return tree
 
 
+def display_workspace_tree(
+    workspace: Dict[str, str], label: str, stderr: bool = False
+) -> 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)
+
+
 PIPELINE_FILE_LINK = "https://s.dvc.org/g/pipeline-files"
 
 
@@ -144,9 +154,9 @@ def init_interactive(
         "Enter the paths for dependencies and outputs of the command.",
         styled=True,
     )
-    if workspace:
-        ui.error_write(build_workspace_tree(workspace), styled=True)
-    ui.error_write(styled=True)
+
+    tree_label = "DVC assumes the following workspace structure:"
+    display_workspace_tree(workspace, tree_label, stderr=True)
     ret.update(
         compact(
             _prompts(prompts, defaults, validator=validator, stream=stream)
@@ -192,13 +202,42 @@ def validate_prompts(
             raise InvalidResponse(msg_format.format(value, suffices[e.errno]))
     elif key in ("code", "data"):
         if not os.path.exists(value):
-            return value, (
-                f"[yellow]'{value}' does not exist in the workspace. "
-                '"exp run" may fail.[/]'
+            typ = "file" if is_file(value) else "directory"
+            return (
+                value,
+                f"[yellow]'{value}' does not exist, "
+                f"the {typ} will be created. ",
             )
     return value
 
 
+def is_file(path: str) -> bool:
+    _, ext = os.path.splitext(path)
+    return bool(ext)
+
+
+def init_deps(deps: Iterable["Dependency"]) -> List[str]:
+    from dvc.dependency import ParamsDependency
+    from dvc.fs.local import localfs
+
+    new_deps = [
+        dep
+        for dep in deps
+        if not isinstance(dep, ParamsDependency) and not dep.exists
+    ]
+    for dep in new_deps:
+        fs_path = dep.fs_path
+        if not is_file(fs_path):
+            localfs.makedirs(fs_path)
+            continue
+
+        localfs.makedirs(localfs.path.parent(fs_path), exist_ok=True)
+        with localfs.open(fs_path, "w", encoding="utf-8"):
+            pass
+
+    return [dep.def_path for dep in new_deps]
+
+
 def init(
     repo: "Repo",
     name: str = "train",
@@ -287,6 +326,10 @@ def init(
         with _disable_logging(), repo.scm_context(autostage=True, quiet=True):
             stage.dump(update_lock=False)
             stage.ignore_outs()
+            if not interactive:
+                label = "Creating experiment project structure:"
+                display_workspace_tree(context, label)
+            init_deps(stage.deps)
             if params:
                 repo.scm_context.track_file(params)
     else:

tests/func/experiments/test_init.py

@@ -44,6 +44,8 @@ def test_init_simple(tmp_dir, scm, dvc, capsys):
             }
         }
     }
+    assert (tmp_dir / "data").read_text() == "data"
+    assert (tmp_dir / "src").is_dir()
 
 
 @pytest.mark.parametrize("interactive", [True, False])
@@ -142,6 +144,8 @@ def test_init_interactive_when_no_path_prompts_need_to_be_asked(
             }
         }
     }
+    assert (tmp_dir / "src").is_dir()
+    assert (tmp_dir / "data").is_dir()
 
 
 def test_when_params_is_omitted_in_interactive_mode(tmp_dir, scm, dvc):
@@ -164,6 +168,8 @@ def test_when_params_is_omitted_in_interactive_mode(tmp_dir, scm, dvc):
         }
     }
     assert not (tmp_dir / "dvc.lock").exists()
+    assert (tmp_dir / "script.py").read_text() == ""
+    assert (tmp_dir / "data").is_dir()
     scm._reset()
     assert scm.is_tracked("dvc.yaml")
     assert not scm.is_tracked("params.yaml")
@@ -194,6 +200,8 @@ def test_init_interactive_params_validation(tmp_dir, dvc, capsys):
             }
         }
     }
+    assert (tmp_dir / "script.py").read_text() == ""
+    assert (tmp_dir / "data").is_dir()
 
     out, err = capsys.readouterr()
     assert (
@@ -209,15 +217,7 @@ def test_init_interactive_params_validation(tmp_dir, dvc, capsys):
 
 
 def test_init_with_no_defaults_interactive(tmp_dir, dvc):
-    inp = io.StringIO(
-        "python script.py\n"
-        "script.py\n"
-        "data\n"
-        "model\n"
-        "n\n"
-        "metric\n"
-        "n\n"
-    )
+    inp = io.StringIO("script.py\n" "data\n" "model\n" "n\n" "metric\n" "n\n")
     init(
         dvc,
         defaults={},
@@ -229,12 +229,14 @@ def test_init_with_no_defaults_interactive(tmp_dir, dvc):
         "stages": {
             "train": {
                 "cmd": "python script.py",
-                "deps": ["python script.py", "script.py"],
+                "deps": ["data", "script.py"],
                 "metrics": [{"metric": {"cache": False}}],
-                "outs": ["data"],
+                "outs": ["model"],
             }
         }
     }
+    assert (tmp_dir / "script.py").read_text() == ""
+    assert (tmp_dir / "data").is_dir()
 
 
 @pytest.mark.parametrize(
@@ -258,9 +260,7 @@ def test_init_with_no_defaults_interactive(tmp_dir, dvc):
     ],
     ids=["non-interactive", "interactive"],
 )
-def test_init_interactive_default(
-    tmp_dir, scm, dvc, interactive, overrides, inp, capsys
-):
+def test_init_default(tmp_dir, scm, dvc, interactive, overrides, inp, capsys):
     (tmp_dir / "params.yaml").dump({"foo": {"bar": 1}})
 
     init(
@@ -284,6 +284,8 @@ def test_init_interactive_default(
         }
     }
     assert not (tmp_dir / "dvc.lock").exists()
+    assert (tmp_dir / "script.py").read_text() == ""
+    assert (tmp_dir / "data").is_dir()
     scm._reset()
     assert scm.is_tracked("dvc.yaml")
     assert scm.is_tracked("params.yaml")
@@ -292,9 +294,11 @@ def test_init_interactive_default(
     out, err = capsys.readouterr()
 
     if interactive:
-        assert "'script.py' does not exist in the workspace." in err
-        assert "'data' does not exist in the workspace." in err
-    assert not out
+        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
+        assert not out
+        return
+    assert "Creating experiment project structure: " in out
 
 
 @pytest.mark.timeout(5, func_only=True)
@@ -365,6 +369,8 @@ def test_init_interactive_live(
         }
     }
     assert not (tmp_dir / "dvc.lock").exists()
+    assert (tmp_dir / "script.py").read_text() == ""
+    assert (tmp_dir / "data").is_dir()
     scm._reset()
     assert scm.is_tracked("dvc.yaml")
     assert scm.is_tracked("params.yaml")
@@ -372,10 +378,13 @@ def test_init_interactive_live(
     assert scm.is_ignored("models")
 
     out, err = capsys.readouterr()
+
     if interactive:
-        assert "'script.py' does not exist in the workspace." in err
-        assert "'data' does not exist in the workspace." in err
-    assert not out
+        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
+        assert not out
+        return
+    assert "Creating experiment project structure: " in out
 
 
 @pytest.mark.parametrize(
@@ -410,6 +419,8 @@ def test_init_with_type_live_and_models_plots_provided(
             }
         }
     }
+    assert (tmp_dir / "src").is_dir()
+    assert (tmp_dir / "data").is_dir()
 
 
 @pytest.mark.parametrize(
@@ -443,3 +454,5 @@ def test_init_with_type_default_and_live_provided(
             }
         }
     }
+    assert (tmp_dir / "src").is_dir()
+    assert (tmp_dir / "data").is_dir()