environment : added machinery to collect modifications to the environment and apply them later

3 files changed, 209 insertions, 1 deletions

diff --git a/lib/spack/spack/environment.py b/lib/spack/spack/environment.py
new file mode 100644
index 0000000000..7d3d7af0de
--- /dev/null
+++ b/lib/spack/spack/environment.py
@@ -0,0 +1,157 @@
+import os
+import os.path
+import collections
+
+
+class SetEnv(object):
+    def __init__(self, name, value, **kwargs):
+        self.name = name
+        self.value = value
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def execute(self):
+        os.environ[self.name] = str(self.value)
+
+
+class UnsetEnv(object):
+    def __init__(self, name, **kwargs):
+        self.name = name
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def execute(self):
+        os.environ.pop(self.name, None)  # Avoid throwing if the variable was not set
+
+
+class AppendPath(object):
+    def __init__(self, name, path, **kwargs):
+        self.name = name
+        self.path = path
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def execute(self):
+        environment_value = os.environ.get(self.name, '')
+        directories = environment_value.split(':') if environment_value else []
+        # TODO : Check if this is a valid directory name
+        directories.append(os.path.normpath(self.path))
+        os.environ[self.name] = ':'.join(directories)
+
+
+class PrependPath(object):
+    def __init__(self, name, path, **kwargs):
+        self.name = name
+        self.path = path
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def execute(self):
+        environment_value = os.environ.get(self.name, '')
+        directories = environment_value.split(':') if environment_value else []
+        # TODO : Check if this is a valid directory name
+        directories = [os.path.normpath(self.path)] + directories
+        os.environ[self.name] = ':'.join(directories)
+
+
+class RemovePath(object):
+    def __init__(self, name, path, **kwargs):
+        self.name = name
+        self.path = path
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def execute(self):
+        environment_value = os.environ.get(self.name, '')
+        directories = environment_value.split(':') if environment_value else []
+        directories = [os.path.normpath(x) for x in directories if x != os.path.normpath(self.path)]
+        os.environ[self.name] = ':'.join(directories)
+
+
+class EnvironmentModifications(object):
+    """
+    Keeps track of requests to modify the current environment
+    """
+
+    def __init__(self):
+        self.env_modifications = []
+
+    def __iter__(self):
+        return iter(self.env_modifications)
+
+    def set_env(self, name, value, **kwargs):
+        """
+        Stores in the current object a request to set an environment variable
+
+        Args:
+            name: name of the environment variable to be set
+            value: value of the environment variable
+        """
+        item = SetEnv(name, value, **kwargs)
+        self.env_modifications.append(item)
+
+    def unset_env(self, name, **kwargs):
+        """
+        Stores in the current object a request to unset an environment variable
+
+        Args:
+            name: name of the environment variable to be set
+        """
+        item = UnsetEnv(name, **kwargs)
+        self.env_modifications.append(item)
+
+    def append_path(self, name, path, **kwargs):
+        """
+        Stores in the current object a request to append a path to a path list
+
+        Args:
+            name: name of the path list in the environment
+            path: path to be appended
+        """
+        item = AppendPath(name, path, **kwargs)
+        self.env_modifications.append(item)
+
+    def prepend_path(self, name, path, **kwargs):
+        """
+        Same as `append_path`, but the path is pre-pended
+
+        Args:
+            name: name of the path list in the environment
+            path: path to be pre-pended
+        """
+        item = PrependPath(name, path, **kwargs)
+        self.env_modifications.append(item)
+
+    def remove_path(self, name, path, **kwargs):
+        """
+        Stores in the current object a request to remove a path from a path list
+
+        Args:
+            name: name of the path list in the environment
+            path: path to be removed
+        """
+        item = RemovePath(name, path, **kwargs)
+        self.env_modifications.append(item)
+
+
+def validate_environment_modifications(env):
+    modifications = collections.defaultdict(list)
+    for item in env:
+        modifications[item.name].append(item)
+    return modifications
+
+
+def apply_environment_modifications(env):
+    """
+    Modifies the current environment according to the request in env
+
+    Args:
+        env: object storing modifications to the environment
+    """
+    modifications = validate_environment_modifications(env)
+
+    # Cycle over the environment variables that will be modified
+    for variable, actions in modifications.items():
+        # Execute all the actions in the order they were issued
+        for x in actions:
+            x.execute()
diff --git a/lib/spack/spack/test/__init__.py b/lib/spack/spack/test/__init__.py
index d5d8b64765..cd842561e6 100644
--- a/lib/spack/spack/test/__init__.py
+++ b/lib/spack/spack/test/__init__.py
@@ -66,7 +66,8 @@ test_names = ['versions',
               'database',
               'namespace_trie',
               'yaml',
-              'sbang']
+              'sbang',
+              'environment']
 
 
 def list_tests():
diff --git a/lib/spack/spack/test/environment.py b/lib/spack/spack/test/environment.py
new file mode 100644
index 0000000000..dff3863d32
--- /dev/null
+++ b/lib/spack/spack/test/environment.py
@@ -0,0 +1,50 @@
+import unittest
+import os
+from spack.environment import EnvironmentModifications, apply_environment_modifications
+
+
+class EnvironmentTest(unittest.TestCase):
+    def setUp(self):
+        os.environ.clear()
+        os.environ['UNSET_ME'] = 'foo'
+        os.environ['EMPTY_PATH_LIST'] = ''
+        os.environ['PATH_LIST'] = '/path/second:/path/third'
+        os.environ['REMOVE_PATH_LIST'] = '/a/b:/duplicate:/a/c:/remove/this:/a/d:/duplicate/:/f/g'
+
+    def test_set_env(self):
+        env = EnvironmentModifications()
+        env.set_env('A', 'dummy value')
+        env.set_env('B', 3)
+        apply_environment_modifications(env)
+        self.assertEqual('dummy value', os.environ['A'])
+        self.assertEqual(str(3), os.environ['B'])
+
+    def test_unset_env(self):
+        env = EnvironmentModifications()
+        self.assertEqual('foo', os.environ['UNSET_ME'])
+        env.unset_env('UNSET_ME')
+        apply_environment_modifications(env)
+        self.assertRaises(KeyError, os.environ.__getitem__, 'UNSET_ME')
+
+    def test_path_manipulation(self):
+        env = EnvironmentModifications()
+
+        env.append_path('PATH_LIST', '/path/last')
+        env.prepend_path('PATH_LIST', '/path/first')
+
+        env.append_path('EMPTY_PATH_LIST', '/path/middle')
+        env.append_path('EMPTY_PATH_LIST', '/path/last')
+        env.prepend_path('EMPTY_PATH_LIST', '/path/first')
+
+        env.append_path('NEWLY_CREATED_PATH_LIST', '/path/middle')
+        env.append_path('NEWLY_CREATED_PATH_LIST', '/path/last')
+        env.prepend_path('NEWLY_CREATED_PATH_LIST', '/path/first')
+
+        env.remove_path('REMOVE_PATH_LIST', '/remove/this')
+        env.remove_path('REMOVE_PATH_LIST', '/duplicate/')
+
+        apply_environment_modifications(env)
+        self.assertEqual('/path/first:/path/second:/path/third:/path/last', os.environ['PATH_LIST'])
+        self.assertEqual('/path/first:/path/middle:/path/last', os.environ['EMPTY_PATH_LIST'])
+        self.assertEqual('/path/first:/path/middle:/path/last', os.environ['NEWLY_CREATED_PATH_LIST'])
+        self.assertEqual('/a/b:/a/c:/a/d:/f/g', os.environ['REMOVE_PATH_LIST'])