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

4 files changed, 128 insertions, 21 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
diff --git a/share/spack/setup-env.csh b/share/spack/setup-env.csh
index 0f67ef7906..a153374028 100755
--- a/share/spack/setup-env.csh
+++ b/share/spack/setup-env.csh
@@ -39,10 +39,8 @@ if ($?SPACK_ROOT) then
     alias spack          'set _sp_args = (\!*); source $_spack_share_dir/csh/spack.csh'
     alias _spack_pathadd 'set _pa_args = (\!*) && source $_spack_share_dir/csh/pathadd.csh'
 
-    # Shamelessly stolen from setup-env.sh
-    set _sp_sys_type    = `$SPACK_ROOT/bin/spack python -c 'print(spack.architecture.sys_type())'`
-    set _sp_dotkit_root = `$SPACK_ROOT/bin/spack python -c "print(spack.util.path.canonicalize_path(spack.config.get('config:module_roots', {}).get('dotkit')))"`
-    set _sp_tcl_root    = `$SPACK_ROOT/bin/spack python -c "print(spack.util.path.canonicalize_path(spack.config.get('config:module_roots', {}).get('tcl')))"`
+    # Set variables needed by this script
+    eval `spack --print-shell-vars csh`
 
     # Set up modules and dotkit search paths in the user environment
     _spack_pathadd DK_NODE    "$_sp_dotkit_root/$_sp_sys_type"
diff --git a/share/spack/setup-env.sh b/share/spack/setup-env.sh
index b983299975..023056d103 100755
--- a/share/spack/setup-env.sh
+++ b/share/spack/setup-env.sh
@@ -218,34 +218,27 @@ if ! _spack_fn_exists use && ! _spack_fn_exists module; then
 	need_module="yes"
 fi;
 
+
 #
-# build and make available environment-modules
+# make available environment-modules
 #
 if [ "${need_module}" = "yes" ]; then
-    #check if environment-modules is installed
-    module_prefix="$(spack location -i "environment-modules" 2>&1 || echo "not_installed")"
-    module_prefix=$(echo "${module_prefix}" | tail -n 1)
-    if [ "${module_prefix}" != "not_installed" ]; then
+    eval `spack --print-shell-vars sh,modules`
+
+    # _sp_module_prefix is set by spack --print-sh-vars
+    if [ "${_sp_module_prefix}" != "not_installed" ]; then
         #activate it!
-        export MODULE_PREFIX=${module_prefix}
+        export MODULE_PREFIX=${_sp_module_prefix}
         _spack_pathadd PATH "${MODULE_PREFIX}/Modules/bin"
         module() { eval `${MODULE_PREFIX}/Modules/bin/modulecmd ${SPACK_SHELL} $*`; }
     fi;
+else
+    eval `spack --print-shell-vars sh`
 fi;
 
 #
-# Set up modules and dotkit search paths in the user environment
+# set module system roots
 #
-
-_python_command=$(printf  "%s\\\n%s\\\n%s" \
-"print(\'_sp_sys_type={0}\'.format(spack.architecture.sys_type()))" \
-"print(\'_sp_dotkit_root={0}\'.format(spack.util.path.canonicalize_path(spack.config.get(\'config:module_roots\', {}).get(\'dotkit\'))))" \
-"print(\'_sp_tcl_root={0}\'.format(spack.util.path.canonicalize_path(spack.config.get(\'config:module_roots\', {}).get(\'tcl\'))))"
-)
-
-_assignment_command=$(spack-python -c "exec('${_python_command}')")
-eval ${_assignment_command}
-
 _spack_pathadd DK_NODE    "${_sp_dotkit_root%/}/$_sp_sys_type"
 _spack_pathadd MODULEPATH "${_sp_tcl_root%/}/$_sp_sys_type"