add `spack --print-shell-vars` to speed up `setup-env.[c]sh` (#8101)

- The setup-env.sh script currently makes two calls to spack, but it
  should only need to make one.

- Add a fast-path shell setup routine in `main.py` to allow the shell
  setup to happen in a single, fast call that doesn't load more than it
  needs to.

- This simplifies setup code, as it has to eval what Spack prints

- TODO: consider eventually making the whole setup script the output of a
  spack command

2 files changed, 116 insertions, 0 deletions

diff --git a/lib/spack/spack/main.py b/lib/spack/spack/main.py
index 6a3f964ab8..08c293b6c0 100644
--- a/lib/spack/spack/main.py
+++ b/lib/spack/spack/main.py
@@ -41,12 +41,15 @@ import llnl.util.tty as tty
 from llnl.util.tty.log import log_output
 
 import spack
+import spack.architecture
 import spack.config
 import spack.cmd
 import spack.hooks
 import spack.paths
 import spack.repo
+import spack.store
 import spack.util.debug
+import spack.util.path
 from spack.error import SpackError
 
 
@@ -362,6 +365,10 @@ def make_argument_parser(**kwargs):
     parser.add_argument(
         '-V', '--version', action='store_true',
         help='show version number and exit')
+    parser.add_argument(
+        '--print-shell-vars', action='store',
+        help="print info needed by setup-env.[c]sh")
+
     return parser
 
 
@@ -547,6 +554,46 @@ def _profile_wrapper(command, parser, args, unknown_args):
         stats.print_stats(nlines)
 
 
+def print_setup_info(*info):
+    """Print basic information needed by setup-env.[c]sh.
+
+    Args:
+        info (list of str): list of things to print: comma-separated list
+            of 'csh', 'sh', or 'modules'
+
+    This is in ``main.py`` to make it fast; the setup scripts need to
+    invoke spack in login scripts, and it needs to be quick.
+
+    """
+    shell = 'csh' if 'csh' in info else 'sh'
+
+    def shell_set(var, value):
+        if shell == 'sh':
+            print("%s='%s'" % (var, value))
+        elif shell == 'csh':
+            print("set %s = '%s'" % (var, value))
+        else:
+            tty.die('shell must be sh or csh')
+
+    # print sys type
+    shell_set('_sp_sys_type', spack.architecture.sys_type())
+
+    # print roots for all module systems
+    module_roots = spack.config.get('config:module_roots')
+    for name, path in module_roots.items():
+        path = spack.util.path.canonicalize_path(path)
+        shell_set('_sp_%s_root' % name, path)
+
+    # print environment module system if available. This can be expensive
+    # on clusters, so skip it if not needed.
+    if 'modules' in info:
+        specs = spack.store.db.query('environment-modules')
+        if specs:
+            shell_set('module_prefix', specs[-1].prefix)
+        else:
+            shell_set('module_prefix', 'not_installed')
+
+
 def main(argv=None):
     """This is the entry point for the Spack command.
 
@@ -562,6 +609,10 @@ def main(argv=None):
     parser.add_argument('command', nargs=argparse.REMAINDER)
     args, unknown = parser.parse_known_args(argv)
 
+    if args.print_shell_vars:
+        print_setup_info(*args.print_shell_vars.split(','))
+        return 0
+
     # Just print help and exit if run with no arguments at all
     no_args = (len(sys.argv) == 1) if argv is None else (len(argv) == 0)
     if no_args:
diff --git a/lib/spack/spack/test/cmd/print_shell_vars.py b/lib/spack/spack/test/cmd/print_shell_vars.py
new file mode 100644
index 0000000000..45019d9178
--- /dev/null
+++ b/lib/spack/spack/test/cmd/print_shell_vars.py
@@ -0,0 +1,65 @@
+##############################################################################
+# Copyright (c) 2013-2018, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+#
+# This file is part of Spack.
+# Created by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
+# LLNL-CODE-647188
+#
+# For details, see https://github.com/spack/spack
+# Please also see the NOTICE and LICENSE files for our notice and the LGPL.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License (as
+# published by the Free Software Foundation) version 2.1, February 1999.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
+# conditions of the GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public
+# License along with this program; if not, write to the Free Software
+# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
+##############################################################################
+from spack.main import print_setup_info
+
+
+def test_print_shell_vars_sh(capsys):
+    print_setup_info('sh')
+    out, _ = capsys.readouterr()
+
+    assert "_sp_sys_type=" in out
+    assert "_sp_tcl_root=" in out
+    assert "_sp_lmod_root=" in out
+    assert "module_prefix" not in out
+
+
+def test_print_shell_vars_csh(capsys):
+    print_setup_info('csh')
+    out, _ = capsys.readouterr()
+
+    assert "set _sp_sys_type = " in out
+    assert "set _sp_tcl_root = " in out
+    assert "set _sp_lmod_root = " in out
+    assert "set module_prefix = " not in out
+
+
+def test_print_shell_vars_sh_modules(capsys):
+    print_setup_info('sh', 'modules')
+    out, _ = capsys.readouterr()
+
+    assert "_sp_sys_type=" in out
+    assert "_sp_tcl_root=" in out
+    assert "_sp_lmod_root=" in out
+    assert "module_prefix=" in out
+
+
+def test_print_shell_vars_csh_modules(capsys):
+    print_setup_info('csh', 'modules')
+    out, _ = capsys.readouterr()
+
+    assert "set _sp_sys_type = " in out
+    assert "set _sp_tcl_root = " in out
+    assert "set _sp_lmod_root = " in out
+    assert "set module_prefix = " in out