Add `--create` to `spack env activate` (#40896)

Add `--create` option to `env activate` to allow users to create and activate in one command.


---------

Co-authored-by: Wouter Deconinck <wdconinc@gmail.com>
Co-authored-by: Tamara Dahlgren <35777542+tldahlgren@users.noreply.github.com>
Co-authored-by: psakievich <psakievich@users.noreply.github.com>

3 files changed, 197 insertions, 107 deletions

diff --git a/lib/spack/docs/environments.rst b/lib/spack/docs/environments.rst
index 63fe483876..93db76b332 100644
--- a/lib/spack/docs/environments.rst
+++ b/lib/spack/docs/environments.rst
@@ -142,6 +142,21 @@ user's prompt to begin with the environment name in brackets.
    $ spack env activate -p myenv
    [myenv] $ ...
 
+The ``activate`` command can also be used to create a new environment, if it is
+not already defined, by adding the ``--create`` flag. Managed and anonymous
+environments, anonymous environments are explained in the next section,
+can both be created using the same flags that `spack env create` accepts.
+If an environment already exists then spack will simply activate it and ignore the
+create specific flags.
+
+.. code-block:: console
+   
+   $ spack env activate --create -p myenv
+   # ...
+   # [creates if myenv does not exist yet]
+   # ...
+   [myenv] $ ...
+
 To deactivate an environment, use the command:
 
 .. code-block:: console
diff --git a/lib/spack/spack/cmd/env.py b/lib/spack/spack/cmd/env.py
index 72af8de9aa..ae5b320b19 100644
--- a/lib/spack/spack/cmd/env.py
+++ b/lib/spack/spack/cmd/env.py
@@ -55,6 +55,104 @@ subcommands = [
 
 
 #
+# env create
+#
+def env_create_setup_parser(subparser):
+    """create a new environment"""
+    subparser.add_argument(
+        "env_name",
+        metavar="env",
+        help=(
+            "name of managed environment or directory of the anonymous env "
+            "(when using --dir/-d) to activate"
+        ),
+    )
+    subparser.add_argument(
+        "-d", "--dir", action="store_true", help="create an environment in a specific directory"
+    )
+    subparser.add_argument(
+        "--keep-relative",
+        action="store_true",
+        help="copy relative develop paths verbatim into the new environment"
+        " when initializing from envfile",
+    )
+    view_opts = subparser.add_mutually_exclusive_group()
+    view_opts.add_argument(
+        "--without-view", action="store_true", help="do not maintain a view for this environment"
+    )
+    view_opts.add_argument(
+        "--with-view",
+        help="specify that this environment should maintain a view at the"
+        " specified path (by default the view is maintained in the"
+        " environment directory)",
+    )
+    subparser.add_argument(
+        "envfile",
+        nargs="?",
+        default=None,
+        help="either a lockfile (must end with '.json' or '.lock') or a manifest file",
+    )
+
+
+def env_create(args):
+    if args.with_view:
+        # Expand relative paths provided on the command line to the current working directory
+        # This way we interpret `spack env create --with-view ./view --dir ./env` as
+        # a view in $PWD/view, not $PWD/env/view. This is different from specifying a relative
+        # path in the manifest, which is resolved relative to the manifest file's location.
+        with_view = os.path.abspath(args.with_view)
+    elif args.without_view:
+        with_view = False
+    else:
+        # Note that 'None' means unspecified, in which case the Environment
+        # object could choose to enable a view by default. False means that
+        # the environment should not include a view.
+        with_view = None
+
+    env = _env_create(
+        args.env_name,
+        init_file=args.envfile,
+        dir=args.dir,
+        with_view=with_view,
+        keep_relative=args.keep_relative,
+    )
+
+    # Generate views, only really useful for environments created from spack.lock files.
+    env.regenerate_views()
+
+
+def _env_create(name_or_path, *, init_file=None, dir=False, with_view=None, keep_relative=False):
+    """Create a new environment, with an optional yaml description.
+
+    Arguments:
+        name_or_path (str): name of the environment to create, or path to it
+        init_file (str or file): optional initialization file -- can be
+            a JSON lockfile (*.lock, *.json) or YAML manifest file
+        dir (bool): if True, create an environment in a directory instead
+            of a named environment
+        keep_relative (bool): if True, develop paths are copied verbatim into
+            the new environment file, otherwise they may be made absolute if the
+            new environment is in a different location
+    """
+    if not dir:
+        env = ev.create(
+            name_or_path, init_file=init_file, with_view=with_view, keep_relative=keep_relative
+        )
+        tty.msg("Created environment '%s' in %s" % (name_or_path, env.path))
+        tty.msg("You can activate this environment with:")
+        tty.msg("  spack env activate %s" % (name_or_path))
+        return env
+
+    env = ev.create_in_dir(
+        name_or_path, init_file=init_file, with_view=with_view, keep_relative=keep_relative
+    )
+    tty.msg("Created environment in %s" % env.path)
+    tty.msg("You can activate this environment with:")
+    tty.msg("  spack env activate %s" % env.path)
+    return env
+
+
+#
 # env activate
 #
 def env_activate_setup_parser(subparser):
@@ -118,22 +216,46 @@ def env_activate_setup_parser(subparser):
         help="decorate the command line prompt when activating",
     )
 
-    env_options = subparser.add_mutually_exclusive_group()
-    env_options.add_argument(
+    subparser.add_argument(
         "--temp",
         action="store_true",
         default=False,
         help="create and activate an environment in a temporary directory",
     )
-    env_options.add_argument(
-        "-d", "--dir", default=None, help="activate the environment in this directory"
+    subparser.add_argument(
+        "--create",
+        action="store_true",
+        default=False,
+        help="create and activate the environment if it doesn't exist",
+    )
+    subparser.add_argument(
+        "--envfile",
+        nargs="?",
+        default=None,
+        help="either a lockfile (must end with '.json' or '.lock') or a manifest file",
+    )
+    subparser.add_argument(
+        "--keep-relative",
+        action="store_true",
+        help="copy relative develop paths verbatim into the new environment"
+        " when initializing from envfile",
     )
-    env_options.add_argument(
+    subparser.add_argument(
+        "-d",
+        "--dir",
+        default=False,
+        action="store_true",
+        help="activate environment based on the directory supplied",
+    )
+    subparser.add_argument(
         metavar="env",
-        dest="activate_env",
+        dest="env_name",
         nargs="?",
         default=None,
-        help="name of environment to activate",
+        help=(
+            "name of managed environment or directory of the anonymous env"
+            " (when using --dir/-d) to activate"
+        ),
     )
 
 
@@ -162,11 +284,17 @@ def env_activate(args):
     if args.env or args.no_env or args.env_dir:
         tty.die("Calling spack env activate with --env, --env-dir and --no-env is ambiguous")
 
-    env_name_or_dir = args.activate_env or args.dir
+    # special parser error handling relative to the --temp flag
+    temp_conflicts = iter([args.keep_relative, args.dir, args.env_name, args.with_view])
+    if args.temp and any(temp_conflicts):
+        tty.die(
+            "spack env activate --temp cannot be combined with managed environments, --with-view,"
+            " --keep-relative, or --dir."
+        )
 
     # When executing `spack env activate` without further arguments, activate
     # the default environment. It's created when it doesn't exist yet.
-    if not env_name_or_dir and not args.temp:
+    if not args.env_name and not args.temp:
         short_name = "default"
         if not ev.exists(short_name):
             ev.create(short_name)
@@ -185,17 +313,25 @@ def env_activate(args):
         _tty_info(f"Created and activated temporary environment in {env_path}")
 
     # Managed environment
-    elif ev.exists(env_name_or_dir) and not args.dir:
-        env_path = ev.root(env_name_or_dir)
-        short_name = env_name_or_dir
+    elif ev.exists(args.env_name) and not args.dir:
+        env_path = ev.root(args.env_name)
+        short_name = args.env_name
 
     # Environment directory
-    elif ev.is_env_dir(env_name_or_dir):
-        env_path = os.path.abspath(env_name_or_dir)
+    elif ev.is_env_dir(args.env_name):
+        env_path = os.path.abspath(args.env_name)
         short_name = os.path.basename(env_path)
 
+    # create if user requested, and then recall recursively
+    elif args.create:
+        tty.set_msg_enabled(False)
+        env_create(args)
+        tty.set_msg_enabled(True)
+        env_activate(args)
+        return
+
     else:
-        tty.die("No such environment: '%s'" % env_name_or_dir)
+        tty.die("No such environment: '%s'" % args.env_name)
 
     env_prompt = "[%s]" % short_name
 
@@ -291,97 +427,6 @@ def env_deactivate(args):
 
 
 #
-# env create
-#
-def env_create_setup_parser(subparser):
-    """create a new environment"""
-    subparser.add_argument("create_env", metavar="env", help="name of environment to create")
-    subparser.add_argument(
-        "-d", "--dir", action="store_true", help="create an environment in a specific directory"
-    )
-    subparser.add_argument(
-        "--keep-relative",
-        action="store_true",
-        help="copy relative develop paths verbatim into the new environment"
-        " when initializing from envfile",
-    )
-    view_opts = subparser.add_mutually_exclusive_group()
-    view_opts.add_argument(
-        "--without-view", action="store_true", help="do not maintain a view for this environment"
-    )
-    view_opts.add_argument(
-        "--with-view",
-        help="specify that this environment should maintain a view at the"
-        " specified path (by default the view is maintained in the"
-        " environment directory)",
-    )
-    subparser.add_argument(
-        "envfile",
-        nargs="?",
-        default=None,
-        help="either a lockfile (must end with '.json' or '.lock') or a manifest file",
-    )
-
-
-def env_create(args):
-    if args.with_view:
-        # Expand relative paths provided on the command line to the current working directory
-        # This way we interpret `spack env create --with-view ./view --dir ./env` as
-        # a view in $PWD/view, not $PWD/env/view. This is different from specifying a relative
-        # path in the manifest, which is resolved relative to the manifest file's location.
-        with_view = os.path.abspath(args.with_view)
-    elif args.without_view:
-        with_view = False
-    else:
-        # Note that 'None' means unspecified, in which case the Environment
-        # object could choose to enable a view by default. False means that
-        # the environment should not include a view.
-        with_view = None
-
-    env = _env_create(
-        args.create_env,
-        init_file=args.envfile,
-        dir=args.dir,
-        with_view=with_view,
-        keep_relative=args.keep_relative,
-    )
-
-    # Generate views, only really useful for environments created from spack.lock files.
-    env.regenerate_views()
-
-
-def _env_create(name_or_path, *, init_file=None, dir=False, with_view=None, keep_relative=False):
-    """Create a new environment, with an optional yaml description.
-
-    Arguments:
-        name_or_path (str): name of the environment to create, or path to it
-        init_file (str or file): optional initialization file -- can be
-            a JSON lockfile (*.lock, *.json) or YAML manifest file
-        dir (bool): if True, create an environment in a directory instead
-            of a named environment
-        keep_relative (bool): if True, develop paths are copied verbatim into
-            the new environment file, otherwise they may be made absolute if the
-            new environment is in a different location
-    """
-    if not dir:
-        env = ev.create(
-            name_or_path, init_file=init_file, with_view=with_view, keep_relative=keep_relative
-        )
-        tty.msg("Created environment '%s' in %s" % (name_or_path, env.path))
-        tty.msg("You can activate this environment with:")
-        tty.msg("  spack env activate %s" % (name_or_path))
-        return env
-
-    env = ev.create_in_dir(
-        name_or_path, init_file=init_file, with_view=with_view, keep_relative=keep_relative
-    )
-    tty.msg("Created environment in %s" % env.path)
-    tty.msg("You can activate this environment with:")
-    tty.msg("  spack env activate %s" % env.path)
-    return env
-
-
-#
 # env remove
 #
 def env_remove_setup_parser(subparser):
diff --git a/lib/spack/spack/test/cmd/env.py b/lib/spack/spack/test/cmd/env.py
index d71a2412b3..aaf2b83e23 100644
--- a/lib/spack/spack/test/cmd/env.py
+++ b/lib/spack/spack/test/cmd/env.py
@@ -2956,7 +2956,9 @@ def test_query_develop_specs(tmpdir):
 )
 def test_activation_and_deactiviation_ambiguities(method, env, no_env, env_dir, capsys):
     """spack [-e x | -E | -D x/]  env [activate | deactivate] y are ambiguous"""
-    args = Namespace(shell="sh", activate_env="a", env=env, no_env=no_env, env_dir=env_dir)
+    args = Namespace(
+        shell="sh", env_name="a", env=env, no_env=no_env, env_dir=env_dir, keep_relative=False
+    )
     with pytest.raises(SystemExit):
         method(args)
     _, err = capsys.readouterr()
@@ -2997,6 +2999,34 @@ def test_activate_temp(monkeypatch, tmpdir):
     assert ev.is_env_dir(str(tmpdir))
 
 
+@pytest.mark.parametrize(
+    "conflict_arg", [["--dir"], ["--keep-relative"], ["--with-view", "foo"], ["env"]]
+)
+def test_activate_parser_conflicts_with_temp(conflict_arg):
+    with pytest.raises(SpackCommandError):
+        env("activate", "--sh", "--temp", *conflict_arg)
+
+
+def test_create_and_activate_managed(tmp_path):
+    with fs.working_dir(str(tmp_path)):
+        shell = env("activate", "--without-view", "--create", "--sh", "foo")
+        active_env_var = next(line for line in shell.splitlines() if ev.spack_env_var in line)
+        assert str(tmp_path) in active_env_var
+        active_ev = ev.active_environment()
+        assert "foo" == active_ev.name
+        env("deactivate")
+
+
+def test_create_and_activate_unmanaged(tmp_path):
+    with fs.working_dir(str(tmp_path)):
+        env_dir = os.path.join(str(tmp_path), "foo")
+        shell = env("activate", "--without-view", "--create", "--sh", "-d", env_dir)
+        active_env_var = next(line for line in shell.splitlines() if ev.spack_env_var in line)
+        assert str(env_dir) in active_env_var
+        assert ev.is_env_dir(env_dir)
+        env("deactivate")
+
+
 def test_activate_default(monkeypatch):
     """Tests whether `spack env activate` creates / activates the default
     environment"""