environment modules : added function to construct them from source files

lib/spack/spack/environment.py

@@ -26,7 +26,9 @@
 import inspect
 import os
 import os.path
-
+import subprocess
+import shlex
+import json
 
 class NameModifier(object):
     def __init__(self, name, **kwargs):
@@ -240,6 +242,78 @@ def apply_modifications(self):
             for x in actions:
                 x.execute()
 
+    @staticmethod
+    def from_sourcing_files(*args, **kwargs):
+        """
+        Creates an instance of EnvironmentModifications that, if executed,
+        has the same effect on the environment as sourcing the files passed as
+        parameters
+
+        Args:
+            *args: list of files to be sourced
+
+        Returns:
+            instance of EnvironmentModifications
+        """
+        env = EnvironmentModifications()
+        # Check if the files are actually there
+        if not all(os.path.isfile(file) for file in args):
+            raise RuntimeError('trying to source non-existing files')
+        # Relevant kwd parameters and formats
+        info = dict(kwargs)
+        info.setdefault('shell', '/bin/bash')
+        info.setdefault('shell_options', '-c')
+        info.setdefault('source_command', 'source')
+        info.setdefault('suppress_output', '&> /dev/null')
+        info.setdefault('concatenate_on_success', '&&')
+
+        shell = '{shell}'.format(**info)
+        shell_options = '{shell_options}'.format(**info)
+        source_file = '{source_command} {file} {concatenate_on_success}'
+        dump_environment = 'python -c "import os, json; print json.dumps(dict(os.environ))"'
+        # Construct the command that will be executed
+        command = [source_file.format(file=file, **info) for file in args]
+        command.append(dump_environment)
+        command = ' '.join(command)
+        command = [
+            shell,
+            shell_options,
+            command
+        ]
+
+        # Try to source all the files,
+        proc = subprocess.Popen(command, stdout=subprocess.PIPE, env=os.environ)
+        proc.wait()
+        if proc.returncode != 0:
+            raise RuntimeError('sourcing files returned a non-zero exit code')
+        output = ''.join([line for line in proc.stdout])
+        # Construct a dictionary with all the variables in the environment
+        after_source_env = dict(json.loads(output))
+
+        # Filter variables that are due to how we source
+        after_source_env.pop('SHLVL')
+        after_source_env.pop('_')
+        after_source_env.pop('PWD')
+
+        # Fill the EnvironmentModifications instance
+        this_environment = dict(os.environ)
+        # New variables
+        new_variables = set(after_source_env) - set(this_environment)
+        for x in new_variables:
+            env.set(x, after_source_env[x])
+        # Variables that have been unset
+        unset_variables = set(this_environment) - set(after_source_env)
+        for x in unset_variables:
+            env.unset(x)
+        # Variables that have been modified
+        common_variables = set(this_environment).intersection(set(after_source_env))
+        modified_variables = [x for x in common_variables if this_environment[x] != after_source_env[x]]
+        for x in modified_variables:
+            # TODO : we may be smarter here, and try to parse if we could compose append_path
+            # TODO : and prepend_path to modify the value
+            env.set(x, after_source_env[x])
+        return env
+
 
 def concatenate_paths(paths, separator=':'):
     """

lib/spack/spack/test/data/sourceme_first.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+export NEW_VAR='new'
+export UNSET_ME='overridden'

lib/spack/spack/test/data/sourceme_second.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+export PATH_LIST='/path/first:/path/second:/path/third:/path/fourth'
+unset EMPTY_PATH_LIST

lib/spack/spack/test/environment.py

@@ -24,7 +24,10 @@
 ##############################################################################
 import unittest
 import os
-from spack.environment import EnvironmentModifications
+
+from spack import spack_root
+from llnl.util.filesystem import join_path
+from spack.environment import EnvironmentModifications, SetEnv, UnsetEnv
 
 
 class EnvironmentTest(unittest.TestCase):
@@ -95,3 +98,37 @@ def test_extend(self):
         self.assertEqual(len(copy_construct), 2)
         for x, y in zip(env, copy_construct):
             assert x is y
+
+    def test_source_files(self):
+        datadir = join_path(spack_root, 'lib', 'spack', 'spack', 'test', 'data')
+        files = [
+            join_path(datadir, 'sourceme_first.sh'),
+            join_path(datadir, 'sourceme_second.sh')
+        ]
+        env = EnvironmentModifications.from_sourcing_files(*files)
+        modifications = env.group_by_name()
+
+        self.assertEqual(len(modifications), 4)
+        # Set new variables
+        self.assertEqual(len(modifications['NEW_VAR']), 1)
+        self.assertTrue(isinstance(modifications['NEW_VAR'][0], SetEnv))
+        self.assertEqual(modifications['NEW_VAR'][0].value, 'new')
+        # Unset variables
+        self.assertEqual(len(modifications['EMPTY_PATH_LIST']), 1)
+        self.assertTrue(isinstance(modifications['EMPTY_PATH_LIST'][0], UnsetEnv))
+        # Modified variables
+        self.assertEqual(len(modifications['UNSET_ME']), 1)
+        self.assertTrue(isinstance(modifications['UNSET_ME'][0], SetEnv))
+        self.assertEqual(modifications['UNSET_ME'][0].value, 'overridden')
+
+        self.assertEqual(len(modifications['PATH_LIST']), 1)
+        self.assertTrue(isinstance(modifications['PATH_LIST'][0], SetEnv))
+        self.assertEqual(modifications['PATH_LIST'][0].value, '/path/first:/path/second:/path/third:/path/fourth')
+
+        # TODO : with reference to the TODO in spack/environment.py
+        # TODO : remove the above and insert
+        # self.assertEqual(len(modifications['PATH_LIST']), 2)
+        # self.assertTrue(isinstance(modifications['PATH_LIST'][0], PrependPath))
+        # self.assertEqual(modifications['PATH_LIST'][0].value, '/path/first')
+        # self.assertTrue(isinstance(modifications['PATH_LIST'][1], AppendPath))
+        # self.assertEqual(modifications['PATH_LIST'][1].value, '/path/fourth')