diff options
-rw-r--r-- | lib/spack/spack/environment.py | 157 | ||||
-rw-r--r-- | lib/spack/spack/test/__init__.py | 3 | ||||
-rw-r--r-- | lib/spack/spack/test/environment.py | 50 |
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']) |