Allow command access to dump/pickle_environment from #8476 (#11434)

* Allow command access to dump/pickle_environment from #8476

2 files changed, 56 insertions, 7 deletions

diff --git a/lib/spack/spack/cmd/build_env.py b/lib/spack/spack/cmd/build_env.py
index 43d74344d9..87961fcb9c 100644
--- a/lib/spack/spack/cmd/build_env.py
+++ b/lib/spack/spack/cmd/build_env.py
@@ -12,8 +12,10 @@ import llnl.util.tty as tty
 import spack.build_environment as build_environment
 import spack.cmd
 import spack.cmd.common.arguments as arguments
+from spack.util.environment import dump_environment, pickle_environment
 
-description = "show install environment for a spec, and run commands"
+description = "run a command in a spec's install environment, " \
+              "or dump its environment to screen or file"
 section = "build"
 level = "long"
 
@@ -21,8 +23,22 @@ level = "long"
 def setup_parser(subparser):
     arguments.add_common_arguments(subparser, ['clean', 'dirty'])
     subparser.add_argument(
+        '--dump', metavar="FILE",
+        help="dump a source-able environment to FILE"
+    )
+    subparser.add_argument(
+        '--pickle', metavar="FILE",
+        help="dump a pickled source-able environment to FILE"
+    )
+    subparser.add_argument(
         'spec', nargs=argparse.REMAINDER,
+        metavar='spec [--] [cmd]...',
         help="specs of package environment to emulate")
+    subparser.epilog\
+        = 'If a command is not specified, the environment will be printed ' \
+        'to standard output (cf /usr/bin/env) unless --dump and/or --pickle ' \
+        'are specified.\n\nIf a command is specified and spec is ' \
+        'multi-word, then the -- separator is obligatory.'
 
 
 def build_env(parser, args):
@@ -49,11 +65,23 @@ def build_env(parser, args):
 
     build_environment.setup_package(spec.package, args.dirty)
 
-    if not cmd:
-        # If no command act like the "env" command and print out env vars.
-        for key, val in os.environ.items():
-            print("%s=%s" % (key, val))
+    if args.dump:
+        # Dump a source-able environment to a text file.
+        tty.msg("Dumping a source-able environment to {0}".format(args.dump))
+        dump_environment(args.dump)
 
-    else:
-        # Otherwise execute the command with the new environment
+    if args.pickle:
+        # Dump a source-able environment to a pickle file.
+        tty.msg(
+            "Pickling a source-able environment to {0}".format(args.pickle))
+        pickle_environment(args.pickle)
+
+    if cmd:
+        # Execute the command with the new environment
         os.execvp(cmd[0], cmd)
+
+    elif not bool(args.pickle or args.dump):
+        # If no command or dump/pickle option act like the "env" command
+        # and print out env vars.
+        for key, val in os.environ.items():
+            print("%s=%s" % (key, val))
diff --git a/lib/spack/spack/test/cmd/build_env.py b/lib/spack/spack/test/cmd/build_env.py
index aaf1e0fbc2..379de521b2 100644
--- a/lib/spack/spack/test/cmd/build_env.py
+++ b/lib/spack/spack/test/cmd/build_env.py
@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+from six.moves import cPickle
 import pytest
 
 from spack.main import SpackCommand, SpackCommandError
@@ -27,3 +28,23 @@ def test_it_just_runs(pkg):
 def test_it_just_fails(pkg, error_cls):
     with pytest.raises(error_cls):
         info(pkg)
+
+
+_out_file = 'env.out'
+
+
+@pytest.mark.usefixtures('config')
+def test_dump(tmpdir):
+    with tmpdir.as_cwd():
+        info('--dump', _out_file, 'zlib')
+        with open(_out_file) as f:
+            assert(any(line.startswith('PATH=') for line in f.readlines()))
+
+
+@pytest.mark.usefixtures('config')
+def test_pickle(tmpdir):
+    with tmpdir.as_cwd():
+        info('--pickle', _out_file, 'zlib')
+        environment = cPickle.load(open(_out_file, 'rb'))
+        assert(type(environment) == dict)
+        assert('PATH' in environment)