From 2911e2bb310017f877ba5e80cedc371b4cf86b1d Mon Sep 17 00:00:00 2001 From: crivella Date: Mon, 14 Oct 2024 13:44:12 +0200 Subject: [PATCH 1/8] Added `wrap_path_env` function --- easybuild/tools/environment.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/easybuild/tools/environment.py b/easybuild/tools/environment.py index cf6bb2e09c..59b0adea3c 100644 --- a/easybuild/tools/environment.py +++ b/easybuild/tools/environment.py @@ -32,6 +32,7 @@ """ import copy import os +import contextlib from easybuild.base import fancylogger from easybuild.tools.build_log import EasyBuildError, dry_run_msg @@ -221,3 +222,35 @@ def sanitize_env(): # unset all $PYTHON* environment variables keys_to_unset = [key for key in os.environ if key.startswith('PYTHON')] unset_env_vars(keys_to_unset, verbose=False) + + +@contextlib.contextmanager +def wrap_path_env(prepend={}, append={}, sep=os.pathsep, strict=False): + """This function is a context manager that temporarily modifies path-like environment variables. + It will prepend and append the values of the given dictionaries to the current environment and restore the + original environment when the context is exited. + """ + orig = {} + for key in prepend.keys() | append.keys(): + if isinstance(sep, dict): + if key not in sep: + if strict: + raise EasyBuildError( + "sep must be a dictionary of strings with keys for all keys in prepend and append" + ) + _sep = os.pathsep + else: + _sep = sep.get(key) + elif isinstance(sep, str): + _sep = sep + else: + raise EasyBuildError("sep must be a string or a dictionary of strings") + val = orig[key] = os.environ.get(key) + path = _sep.join(filter(None, [prepend.get(key, None), val, append.get(key, None)])) + setvar(key, path) + + try: + yield + finally: + for key in orig: + setvar(key, orig[key]) From 6972f22ba03c7bd386cb8113e9cfd1c464896c90 Mon Sep 17 00:00:00 2001 From: crivella Date: Mon, 14 Oct 2024 13:44:17 +0200 Subject: [PATCH 2/8] Added tests --- test/framework/environment.py | 50 +++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/test/framework/environment.py b/test/framework/environment.py index 75919f6a13..51603bed06 100644 --- a/test/framework/environment.py +++ b/test/framework/environment.py @@ -33,6 +33,7 @@ from unittest import TextTestRunner import easybuild.tools.environment as env +from easybuild.tools.build_log import EasyBuildError class EnvironmentTest(EnhancedTestCase): @@ -161,6 +162,55 @@ def test_sanitize_env(self): self.assertEqual(os.getenv('LD_PRELOAD'), None) + def test_wrap_env(self): + """Test wrap_env function.""" + os.environ['TEST_VAR_1'] = '/bar:/foo' + os.environ['TEST_VAR_2'] = '/bar' + os.environ['TEST_VAR_3'] = '/foo' + + prep = { + 'TEST_VAR_1': '/usr/bin:/usr/sbin', + 'TEST_VAR_2': '/usr/bin', + } + appd = { + 'TEST_VAR_1': '/usr/local/bin', + 'TEST_VAR_3': '/usr/local/sbin', + } + seps = { + 'TEST_VAR_3': ';' + } + + # Test prepend and append + self.assertEqual(os.getenv('TEST_VAR_1'), '/bar:/foo') + self.assertEqual(os.getenv('TEST_VAR_2'), '/bar') + self.assertEqual(os.getenv('TEST_VAR_3'), '/foo') + with env.wrap_path_env(prep, appd, sep=seps): + self.assertEqual(os.getenv('TEST_VAR_1'), '/usr/bin:/usr/sbin:/bar:/foo:/usr/local/bin') + self.assertEqual(os.getenv('TEST_VAR_2'), '/usr/bin:/bar') + self.assertEqual(os.getenv('TEST_VAR_3'), '/foo;/usr/local/sbin') + # Test modifying the environment inside the context + os.environ['TEST_VAR_1'] = '' + os.environ['TEST_VAR_2'] = '' + os.environ['TEST_VAR_3'] = '' + self.assertEqual(os.getenv('TEST_VAR_1'), '') + self.assertEqual(os.getenv('TEST_VAR_2'), '') + self.assertEqual(os.getenv('TEST_VAR_3'), '') + self.assertEqual(os.getenv('TEST_VAR_1'), '/bar:/foo') + self.assertEqual(os.getenv('TEST_VAR_2'), '/bar') + self.assertEqual(os.getenv('TEST_VAR_3'), '/foo') + + # Test sep with strict=True + def foo(): + with env.wrap_path_env(prep, appd, sep={}, strict=True): + pass + self.assertErrorRegex(EasyBuildError, "sep must be a .*", foo) + + # Test invalid value for sep + def foo(): + with env.wrap_path_env(prep, appd, sep=None): + pass + self.assertErrorRegex(EasyBuildError, "sep must be a .*", foo) + def suite(): """ returns all the testcases in this module """ From 4ec0b54e24db267fd1d7deea415de028fb6d8469 Mon Sep 17 00:00:00 2001 From: crivella Date: Mon, 14 Oct 2024 14:13:00 +0200 Subject: [PATCH 3/8] lint --- easybuild/tools/environment.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/easybuild/tools/environment.py b/easybuild/tools/environment.py index 59b0adea3c..104ab08723 100644 --- a/easybuild/tools/environment.py +++ b/easybuild/tools/environment.py @@ -225,11 +225,20 @@ def sanitize_env(): @contextlib.contextmanager -def wrap_path_env(prepend={}, append={}, sep=os.pathsep, strict=False): +def wrap_path_env(prepend=None, append=None, sep=os.pathsep, strict=False): """This function is a context manager that temporarily modifies path-like environment variables. It will prepend and append the values of the given dictionaries to the current environment and restore the original environment when the context is exited. + + A custom separator can be specified for each variable by passing a dictionary of strings with the same keys as + prepend and append. + If a key is not present in the sep dictionary, os.pathsep will be used unless strict is True, then an error + will be raised. """ + if prepend is None: + prepend = {} + if append is None: + append = {} orig = {} for key in prepend.keys() | append.keys(): if isinstance(sep, dict): @@ -252,5 +261,4 @@ def wrap_path_env(prepend={}, append={}, sep=os.pathsep, strict=False): try: yield finally: - for key in orig: - setvar(key, orig[key]) + restore_env_vars(orig) From f3f0e4a987c5cf47e81460da7672a38b3eb2ceda Mon Sep 17 00:00:00 2001 From: crivella Date: Mon, 14 Oct 2024 15:24:13 +0200 Subject: [PATCH 4/8] Python2 fix --- easybuild/tools/environment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/easybuild/tools/environment.py b/easybuild/tools/environment.py index 104ab08723..f7dbc9394f 100644 --- a/easybuild/tools/environment.py +++ b/easybuild/tools/environment.py @@ -240,7 +240,7 @@ def wrap_path_env(prepend=None, append=None, sep=os.pathsep, strict=False): if append is None: append = {} orig = {} - for key in prepend.keys() | append.keys(): + for key in set(prepend.keys()) | set(append.keys()): if isinstance(sep, dict): if key not in sep: if strict: From 61a34a58b3087bc3801efce537f05f5b358e9247 Mon Sep 17 00:00:00 2001 From: crivella Date: Mon, 18 Nov 2024 09:49:35 +0100 Subject: [PATCH 5/8] Added `override` and docstrings --- easybuild/tools/environment.py | 36 ++++++++++++++++++++++++++++------ 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/easybuild/tools/environment.py b/easybuild/tools/environment.py index f7dbc9394f..81084de72f 100644 --- a/easybuild/tools/environment.py +++ b/easybuild/tools/environment.py @@ -225,22 +225,46 @@ def sanitize_env(): @contextlib.contextmanager -def wrap_path_env(prepend=None, append=None, sep=os.pathsep, strict=False): - """This function is a context manager that temporarily modifies path-like environment variables. - It will prepend and append the values of the given dictionaries to the current environment and restore the +def wrap_path_env(override=None, prepend=None, append=None, sep=os.pathsep, strict=False): + """This function is a context manager that temporarily modifies environment variables. + It will override or prepend/append the values of the given dictionaries to the current environment and restore the original environment when the context is exited. - A custom separator can be specified for each variable by passing a dictionary of strings with the same keys as - prepend and append. + For path-like variables, a custom separator can be specified for each variable by passing a dictionary of strings + with the same keys as prepend and append. If a key is not present in the sep dictionary, os.pathsep will be used unless strict is True, then an error will be raised. + + Args: + override: A dictionary of environment variables to override. + prepend: A dictionary of environment variables to prepend to. + append: A dictionary of environment variables to append to. + sep: A string or a dictionary of strings to use as separator for each variable. + strict: If True, raise an error if a key is not present in the sep dictionary. """ if prepend is None: prepend = {} if append is None: append = {} + if override is None: + override = {} + + path_keys = set(prepend.keys()) | set(append.keys()) + over_keys = set(override.keys()) + + duplicates = path_keys & over_keys + if duplicates: + raise EasyBuildError( + "The keys in override must not overlap with the keys in prepend or append: '%s'", + " ".join(duplicates) + ) + orig = {} - for key in set(prepend.keys()) | set(append.keys()): + for key in over_keys + orig[key] = os.environ.get(key) + setvar(key, override[key]) + + for key in path_keys: if isinstance(sep, dict): if key not in sep: if strict: From 642432214bba170bcc1ef172310725b514e95cad Mon Sep 17 00:00:00 2001 From: crivella Date: Mon, 18 Nov 2024 09:50:59 +0100 Subject: [PATCH 6/8] lint --- easybuild/tools/environment.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/easybuild/tools/environment.py b/easybuild/tools/environment.py index 81084de72f..74c7e9a2f5 100644 --- a/easybuild/tools/environment.py +++ b/easybuild/tools/environment.py @@ -258,9 +258,9 @@ def wrap_path_env(override=None, prepend=None, append=None, sep=os.pathsep, stri "The keys in override must not overlap with the keys in prepend or append: '%s'", " ".join(duplicates) ) - + orig = {} - for key in over_keys + for key in over_keys: orig[key] = os.environ.get(key) setvar(key, override[key]) From d5679e3b492cbb06cb89c75d7f32c34f6db1a5db Mon Sep 17 00:00:00 2001 From: crivella Date: Mon, 18 Nov 2024 10:32:02 +0100 Subject: [PATCH 7/8] Renamed function adn adjusted tests --- easybuild/tools/environment.py | 2 +- test/framework/environment.py | 70 +++++++++++++++++++++++++--------- 2 files changed, 53 insertions(+), 19 deletions(-) diff --git a/easybuild/tools/environment.py b/easybuild/tools/environment.py index 74c7e9a2f5..6550fb0391 100644 --- a/easybuild/tools/environment.py +++ b/easybuild/tools/environment.py @@ -225,7 +225,7 @@ def sanitize_env(): @contextlib.contextmanager -def wrap_path_env(override=None, prepend=None, append=None, sep=os.pathsep, strict=False): +def wrap_env(override=None, prepend=None, append=None, sep=os.pathsep, strict=False): """This function is a context manager that temporarily modifies environment variables. It will override or prepend/append the values of the given dictionaries to the current environment and restore the original environment when the context is exited. diff --git a/test/framework/environment.py b/test/framework/environment.py index 51603bed06..fa01005eed 100644 --- a/test/framework/environment.py +++ b/test/framework/environment.py @@ -164,9 +164,20 @@ def test_sanitize_env(self): def test_wrap_env(self): """Test wrap_env function.""" - os.environ['TEST_VAR_1'] = '/bar:/foo' - os.environ['TEST_VAR_2'] = '/bar' - os.environ['TEST_VAR_3'] = '/foo' + def reset_env(): + os.environ['TEST_VAR_1'] = '/bar:/foo' + os.environ['TEST_VAR_2'] = '/bar' + os.environ['TEST_VAR_3'] = '/foo' + def check_env(): + self.assertEqual(os.getenv('TEST_VAR_1'), '/bar:/foo') + self.assertEqual(os.getenv('TEST_VAR_2'), '/bar') + self.assertEqual(os.getenv('TEST_VAR_3'), '/foo') + def null_and_check(vars): + for var in vars: + os.environ[var] = '' + for var in vars: + self.assertEqual(os.getenv(var), '') + reset_env() prep = { 'TEST_VAR_1': '/usr/bin:/usr/sbin', @@ -176,41 +187,64 @@ def test_wrap_env(self): 'TEST_VAR_1': '/usr/local/bin', 'TEST_VAR_3': '/usr/local/sbin', } + over = { + 'TEST_VAR_3': 'overridden', + } seps = { 'TEST_VAR_3': ';' } # Test prepend and append - self.assertEqual(os.getenv('TEST_VAR_1'), '/bar:/foo') - self.assertEqual(os.getenv('TEST_VAR_2'), '/bar') - self.assertEqual(os.getenv('TEST_VAR_3'), '/foo') - with env.wrap_path_env(prep, appd, sep=seps): + reset_env() + check_env() + with env.wrap_env(prepend=prep, append=appd, sep=seps): self.assertEqual(os.getenv('TEST_VAR_1'), '/usr/bin:/usr/sbin:/bar:/foo:/usr/local/bin') self.assertEqual(os.getenv('TEST_VAR_2'), '/usr/bin:/bar') self.assertEqual(os.getenv('TEST_VAR_3'), '/foo;/usr/local/sbin') # Test modifying the environment inside the context - os.environ['TEST_VAR_1'] = '' - os.environ['TEST_VAR_2'] = '' - os.environ['TEST_VAR_3'] = '' - self.assertEqual(os.getenv('TEST_VAR_1'), '') - self.assertEqual(os.getenv('TEST_VAR_2'), '') - self.assertEqual(os.getenv('TEST_VAR_3'), '') - self.assertEqual(os.getenv('TEST_VAR_1'), '/bar:/foo') - self.assertEqual(os.getenv('TEST_VAR_2'), '/bar') - self.assertEqual(os.getenv('TEST_VAR_3'), '/foo') + null_and_check(['TEST_VAR_1', 'TEST_VAR_2', 'TEST_VAR_3']) + check_env() # Test sep with strict=True def foo(): - with env.wrap_path_env(prep, appd, sep={}, strict=True): + with env.wrap_env(prepend=prep, append=appd, sep={}, strict=True): pass self.assertErrorRegex(EasyBuildError, "sep must be a .*", foo) # Test invalid value for sep def foo(): - with env.wrap_path_env(prep, appd, sep=None): + with env.wrap_env(prepend=prep, append=appd, sep=None): pass self.assertErrorRegex(EasyBuildError, "sep must be a .*", foo) + # Test override + check_env() + with env.wrap_env(override=prep): + self.assertEqual(os.getenv('TEST_VAR_1'), '/usr/bin:/usr/sbin') + self.assertEqual(os.getenv('TEST_VAR_2'), '/usr/bin') + self.assertEqual(os.getenv('TEST_VAR_3'), '/foo') + # Test modifying the environment inside the context + null_and_check(['TEST_VAR_1', 'TEST_VAR_2']) + check_env() + + # Test override with prepend + with env.wrap_env(override=over, prepend=prep): + self.assertEqual(os.getenv('TEST_VAR_1'), '/usr/bin:/usr/sbin:/bar:/foo') + self.assertEqual(os.getenv('TEST_VAR_2'), '/usr/bin:/bar') + self.assertEqual(os.getenv('TEST_VAR_3'), 'overridden') + null_and_check(['TEST_VAR_1', 'TEST_VAR_2', 'TEST_VAR_3']) + check_env() + + # Test override duplicate key with prepend + def foo(): + with env.wrap_env(override=prep, prepend=prep, sep=None): + pass + self.assertErrorRegex( + EasyBuildError, + "The keys in override must not overlap with the keys in prepend or append.*", + foo + ) + def suite(): """ returns all the testcases in this module """ From 2906508ae580cfdcbd8b8cbe42bde46046be6e00 Mon Sep 17 00:00:00 2001 From: crivella Date: Mon, 18 Nov 2024 10:33:33 +0100 Subject: [PATCH 8/8] lint --- test/framework/environment.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/framework/environment.py b/test/framework/environment.py index fa01005eed..e687772025 100644 --- a/test/framework/environment.py +++ b/test/framework/environment.py @@ -164,20 +164,22 @@ def test_sanitize_env(self): def test_wrap_env(self): """Test wrap_env function.""" + def reset_env(): os.environ['TEST_VAR_1'] = '/bar:/foo' os.environ['TEST_VAR_2'] = '/bar' os.environ['TEST_VAR_3'] = '/foo' + def check_env(): self.assertEqual(os.getenv('TEST_VAR_1'), '/bar:/foo') self.assertEqual(os.getenv('TEST_VAR_2'), '/bar') self.assertEqual(os.getenv('TEST_VAR_3'), '/foo') + def null_and_check(vars): for var in vars: os.environ[var] = '' for var in vars: self.assertEqual(os.getenv(var), '') - reset_env() prep = { 'TEST_VAR_1': '/usr/bin:/usr/sbin',