module_cmd.py: use posix awk; fix stderr redirection (#29440)

Emulates `env -0` in a posix compliant way, avoiding a slow python process, speeds up setting up the build env when modules should load.

2 files changed, 34 insertions, 62 deletions

diff --git a/lib/spack/spack/test/module_parsing.py b/lib/spack/spack/test/module_parsing.py
index 7c7c7066f0..cf6508ca70 100644
--- a/lib/spack/spack/test/module_parsing.py
+++ b/lib/spack/spack/test/module_parsing.py
@@ -8,6 +8,7 @@ import os
 import pytest
 
 import spack
+import spack.util.module_cmd
 from spack.util.module_cmd import (
     get_path_args_from_module_line,
     get_path_from_module_contents,
@@ -21,30 +22,26 @@ test_module_lines = ['prepend-path LD_LIBRARY_PATH /path/to/lib',
                      'setenv LDFLAGS -L/path/to/lib',
                      'prepend-path PATH /path/to/bin']
 
-_test_template = "'. %s 2>&1' % args[1]"
 
-
-def test_module_function_change_env(tmpdir, working_env, monkeypatch):
-    monkeypatch.setattr(spack.util.module_cmd, '_cmd_template', _test_template)
+def test_module_function_change_env(tmpdir, working_env):
     src_file = str(tmpdir.join('src_me'))
     with open(src_file, 'w') as f:
         f.write('export TEST_MODULE_ENV_VAR=TEST_SUCCESS\n')
 
     os.environ['NOT_AFFECTED'] = "NOT_AFFECTED"
-    module('load', src_file)
+    module('load', src_file, module_template='. {0} 2>&1'.format(src_file))
 
     assert os.environ['TEST_MODULE_ENV_VAR'] == 'TEST_SUCCESS'
     assert os.environ['NOT_AFFECTED'] == "NOT_AFFECTED"
 
 
-def test_module_function_no_change(tmpdir, monkeypatch):
-    monkeypatch.setattr(spack.util.module_cmd, '_cmd_template', _test_template)
+def test_module_function_no_change(tmpdir):
     src_file = str(tmpdir.join('src_me'))
     with open(src_file, 'w') as f:
         f.write('echo TEST_MODULE_FUNCTION_PRINT')
 
     old_env = os.environ.copy()
-    text = module('show', src_file)
+    text = module('show', src_file, module_template='. {0} 2>&1'.format(src_file))
 
     assert text == 'TEST_MODULE_FUNCTION_PRINT\n'
     assert os.environ == old_env
diff --git a/lib/spack/spack/util/module_cmd.py b/lib/spack/spack/util/module_cmd.py
index 4e897e85d8..82ce6c07a4 100644
--- a/lib/spack/spack/util/module_cmd.py
+++ b/lib/spack/spack/util/module_cmd.py
@@ -7,7 +7,6 @@
 This module contains routines related to the module command for accessing and
 parsing environment modules.
 """
-import json
 import os
 import re
 import subprocess
@@ -15,70 +14,46 @@ import sys
 
 import llnl.util.tty as tty
 
-import spack
-
 # This list is not exhaustive. Currently we only use load and unload
 # If we need another option that changes the environment, add it here.
 module_change_commands = ['load', 'swap', 'unload', 'purge', 'use', 'unuse']
-py_cmd = 'import os;import json;print(json.dumps(dict(os.environ)))'
-_cmd_template = "'module ' + ' '.join(args) + ' 2>&1'"
-
-
-def module(*args):
-    module_cmd = eval(_cmd_template)  # So we can monkeypatch for testing
-    if args[0] in module_change_commands:
-        # Do the module manipulation, then output the environment in JSON
-        # and read the JSON back in the parent process to update os.environ
-        # For python, we use the same python running the Spack process, because
-        # we can guarantee its existence. We have to do some LD_LIBRARY_PATH
-        # shenanigans to ensure python will run.
-
-        # LD_LIBRARY_PATH under which Spack ran
-        os.environ['SPACK_LD_LIBRARY_PATH'] = spack.main.spack_ld_library_path
-
-        # suppress output from module function
-        module_cmd += ' >/dev/null;'
 
-        # Capture the new LD_LIBRARY_PATH after `module` was run
-        module_cmd += 'export SPACK_NEW_LD_LIBRARY_PATH="$LD_LIBRARY_PATH";'
+# This awk script is a posix alternative to `env -0`
+awk_cmd = (r"""awk 'BEGIN{for(name in ENVIRON)"""
+           r"""printf("%s=%s%c", name, ENVIRON[name], 0)}'""")
 
-        # Set LD_LIBRARY_PATH to value at Spack startup time to ensure that
-        # python executable finds its libraries
-        module_cmd += 'LD_LIBRARY_PATH="$SPACK_LD_LIBRARY_PATH" '
 
-        # Execute the python command
-        module_cmd += '%s -E -c "%s";' % (sys.executable, py_cmd)
+def module(*args, **kwargs):
+    module_cmd = kwargs.get('module_template', 'module ' + ' '.join(args))
 
-        # If LD_LIBRARY_PATH was set after `module`, dump the old value because
-        # we have since corrupted it to ensure python would run.
-        # dump SPACKIGNORE as a placeholder for parsing if LD_LIBRARY_PATH null
-        module_cmd += 'echo "${SPACK_NEW_LD_LIBRARY_PATH:-SPACKIGNORE}"'
-
-        module_p  = subprocess.Popen(module_cmd,
-                                     stdout=subprocess.PIPE,
-                                     stderr=subprocess.STDOUT,
-                                     shell=True,
-                                     executable="/bin/bash")
-
-        # Cray modules spit out warnings that we cannot supress.
-        # This hack skips to the last output (the environment)
-        env_out = str(module_p.communicate()[0].decode()).strip().split('\n')
-
-        # The environment dumped as json
-        env_json = env_out[-2]
-        # Either the uncorrupted $LD_LIBRARY_PATH or SPACKIGNORE
-        new_ld_library_path = env_out[-1]
+    if args[0] in module_change_commands:
+        # Suppress module output
+        module_cmd += r' >/dev/null 2>&1; ' + awk_cmd
+        module_p = subprocess.Popen(
+            module_cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            shell=True,
+            executable="/bin/bash")
+
+        # In Python 3, keys and values of `environ` are byte strings.
+        environ = {}
+        output = module_p.communicate()[0]
+
+        # Loop over each environment variable key=value byte string
+        for entry in output.strip(b'\0').split(b'\0'):
+            # Split variable name and value
+            parts = entry.split(b'=', 1)
+            if len(parts) != 2:
+                continue
+            environ[parts[0]] = parts[1]
 
         # Update os.environ with new dict
-        env_dict = json.loads(env_json)
         os.environ.clear()
-        os.environ.update(env_dict)
-
-        # Override restored LD_LIBRARY_PATH with pre-python value
-        if new_ld_library_path == 'SPACKIGNORE':
-            os.environ.pop('LD_LIBRARY_PATH', None)
+        if sys.version_info >= (3, 2):
+            os.environb.update(environ)  # novermin
         else:
-            os.environ['LD_LIBRARY_PATH'] = new_ld_library_path
+            os.environ.update(environ)
 
     else:
         # Simply execute commands that don't change state and return output