Skip to content

add support for environment variables in set_environment of module generators #4855

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
boegel merged 14 commits into from
May 21, 2025
Merged

add support for environment variables in set_environment of module generators #4855

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
d5fef2c
replace PATH_JOIN_TEMPLATE in ModuleGeneratorLua with _path_join_cmd(…
lexming Apr 24, 2025
8e46ad8
add support for environment variables in set_environment method of mo…
lexming Apr 24, 2025
5f3b4d1
fix formatting of inline comments in test_env
lexming Apr 24, 2025
788e925
update test_make_module_req to adapt to new _path_join_cmd method in …
lexming Apr 24, 2025
510f517
update test_make_module_step to adapt to new _path_join_cmd method in…
lexming Apr 24, 2025
67aeaa2
update test_toy_module_fulltxt to adapt to new _path_join_cmd method …
lexming Apr 24, 2025
72cdac7
update test_toy_python to adapt to new _path_join_cmd method in Modul…
lexming Apr 24, 2025
5c594ad
fix code formatting in test_make_module_step
lexming Apr 24, 2025
ad20bd0
update test_make_module_extra to adapt to new _path_join_cmd method i…
lexming Apr 24, 2025
bb0a777
add 'shell_vars' option to modextravars to control resolution of shel…
lexming May 16, 2025
510e8f0
minor changes to ModuleGenerator w.r.t. resolve environment variables…
boegel May 21, 2025
b3842b4
fix docstring + deprecation warning for ModuleGenerator.unpack_setenv…
boegel May 21, 2025
db919e4
rename 'shell_vars' options in modextravars to 'resolve_env_vars' (+ …
boegel May 21, 2025
3ce7693
Merge pull request #10 from boegel/modextra-envars
lexming May 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 72 additions & 30 deletions easybuild/tools/module_generator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,6 @@
from easybuild.tools.modules import ROOT_ENV_VAR_NAME_PREFIX, EnvironmentModulesC, Lmod, modules_tool
from easybuild.tools.utilities import get_subclasses, nub, quote_str


_log = fancylogger.getLogger('module_generator', fname=False)


Expand Down Expand Up @@ -133,6 +132,15 @@ class ModuleGenerator:
# a single level of indentation
INDENTATION = ' ' * 4

# shell environment variable name: ${__}VAR_NAME_00_SUFFIX
REGEX_SHELL_VAR_PATTERN = r'[A-Z_]+[A-Z0-9_]+'
REGEX_SHELL_VAR = re.compile(rf'\$({REGEX_SHELL_VAR_PATTERN})')
REGEX_QUOTE_SHELL_VAR = re.compile(rf'[\"\']\$({REGEX_SHELL_VAR_PATTERN})[\"\']')

# default options for modextravars
DEFAULT_MODEXTRAVARS_USE_PUSHENV = False
DEFAULT_MODEXTRAVARS_RESOLVE_ENV_VARS = True

def __init__(self, application, fake=False):
"""ModuleGenerator constructor."""
self.app = application
Expand Down Expand Up @@ -422,28 +430,38 @@ def det_installdir(self, modfile):

return res

def unpack_setenv_value(self, env_var_name, env_var_val):
def unpack_setenv_value(self, *args, **kwargs):
"""
DEPRECATED method, should not be used.
Replaced with (internal) _unpack_setenv_value method.
"""
self.log.deprecated("unpack_setenv_value should not be used directly (replaced by internal method)", '6.0')
value, use_pushenv, _ = self._unpack_setenv_value(*args, **kwargs)
return value, use_pushenv

def _unpack_setenv_value(self, env_var_name, env_var_val):
"""
Unpack value that specifies how to define an environment variable with specified name.
"""
use_pushenv = False
use_pushenv = self.DEFAULT_MODEXTRAVARS_USE_PUSHENV
resolve_env_vars = self.DEFAULT_MODEXTRAVARS_RESOLVE_ENV_VARS

# value may be specified as a string, or as a dict for special cases
if isinstance(env_var_val, str):
value = env_var_val

elif isinstance(env_var_val, dict):
use_pushenv = env_var_val.get('pushenv', False)
use_pushenv = env_var_val.get('pushenv', self.DEFAULT_MODEXTRAVARS_USE_PUSHENV)
resolve_env_vars = env_var_val.get('resolve_env_vars', self.DEFAULT_MODEXTRAVARS_RESOLVE_ENV_VARS)
try:
value = env_var_val['value']
except KeyError:
except KeyError as err:
raise EasyBuildError("Required key 'value' is missing in dict that specifies how to set $%s: %s",
env_var_name, env_var_val)
env_var_name, env_var_val) from err
else:
raise EasyBuildError("Incorrect value type for setting $%s environment variable (%s): %s",
env_var_name, type(env_var_val), env_var_val)

return value, use_pushenv
return value, use_pushenv, resolve_env_vars

# From this point on just not implemented methods

Expand Down Expand Up @@ -1056,19 +1074,19 @@ def set_environment(self, key, value, relpath=False):
self.log.info("Not including statement to define environment variable $%s, as specified", key)
return ''

value, use_pushenv = self.unpack_setenv_value(key, value)
set_value, use_pushenv, resolve_env_vars = self._unpack_setenv_value(key, value)

# quotes are needed, to ensure smooth working of EBDEVEL* modulefiles
if relpath:
if value:
val = quote_str(os.path.join('$root', value), tcl=True)
else:
val = '"$root"'
else:
val = quote_str(value, tcl=True)
set_value = os.path.join('$root', set_value) if set_value else '$root'

if resolve_env_vars:
set_value = self.REGEX_SHELL_VAR.sub(r'$::env(\1)', set_value)

# quotes are needed, to ensure smooth working of EBDEVEL* modulefiles
set_value = quote_str(set_value, tcl=True)

env_setter = 'pushenv' if use_pushenv else 'setenv'
return '%s\t%s\t\t%s\n' % (env_setter, key, val)
return f'{env_setter}\t{key}\t\t{set_value}\n'

def swap_module(self, mod_name_out, mod_name_in, guarded=True):
"""
Expand Down Expand Up @@ -1152,12 +1170,14 @@ class ModuleGeneratorLua(ModuleGenerator):
LOAD_TEMPLATE_DEPENDS_ON = 'depends_on("%(mod_name)s")'
IS_LOADED_TEMPLATE = 'isloaded("%s")'

OS_GETENV_TEMPLATE = r'os.getenv("%s")'
PATH_JOIN_TEMPLATE = 'pathJoin(root, "%s")'
UPDATE_PATH_TEMPLATE = '%s_path("%s", %s)'
UPDATE_PATH_TEMPLATE_DELIM = '%s_path("%s", %s, "%s")'

START_STR = '[==['
END_STR = ']==]'
CONCAT_STR = ' .. '

def __init__(self, *args, **kwargs):
"""ModuleGeneratorLua constructor."""
Expand All @@ -1167,6 +1187,20 @@ def __init__(self, *args, **kwargs):
if self.modules_tool.version and LooseVersion(self.modules_tool.version) >= LooseVersion('7.7.38'):
self.DOT_MODULERC = '.modulerc.lua'

@staticmethod
def _path_join_cmd(path):
"Return 'pathJoin' command for given path string"
path_components = [quote_str(p) for p in path.split(os.path.sep) if p]

path_root = quote_str(os.path.sep) if os.path.isabs(path) else 'root'
path_components.insert(0, path_root)

if len(path_components) > 1:
return 'pathJoin(' + ', '.join(path_components) + ')'

# no need for a pathJoin for single component paths
return path_components[0]

def check_version(self, minimal_version_maj, minimal_version_min, minimal_version_patch='0'):
"""
Check the minimal version of the moduletool in the module file
Expand Down Expand Up @@ -1292,10 +1326,9 @@ def getenv_cmd(self, envvar, default=None):
"""
Return module-syntax specific code to get value of specific environment variable.
"""
if default is None:
cmd = 'os.getenv("%s")' % envvar
else:
cmd = 'os.getenv("%s") or "%s"' % (envvar, default)
cmd = self.OS_GETENV_TEMPLATE % envvar
if default is not None:
cmd += f' or "{default}"'
return cmd

def load_module(self, mod_name, recursive_unload=None, depends_on=None, unload_modules=None, multi_dep_mods=None):
Expand Down Expand Up @@ -1448,7 +1481,7 @@ def update_paths(self, key, paths, prepend=True, allow_abs=False, expand_relpath
# use pathJoin for (non-empty) relative paths
if path:
if expand_relpaths:
abspaths.append(self.PATH_JOIN_TEMPLATE % path)
abspaths.append(self._path_join_cmd(path))
else:
abspaths.append(quote_str(path))
else:
Expand Down Expand Up @@ -1513,19 +1546,28 @@ def set_environment(self, key, value, relpath=False):
self.log.info("Not including statement to define environment variable $%s, as specified", key)
return ''

value, use_pushenv = self.unpack_setenv_value(key, value)
set_value, use_pushenv, resolve_env_vars = self._unpack_setenv_value(key, value)

if relpath:
if value:
val = self.PATH_JOIN_TEMPLATE % value
else:
val = 'root'
set_value = self._path_join_cmd(set_value)
if resolve_env_vars:
# replace quoted substring with env var with os.getenv statement
# example: pathJoin(root, "$HOME") -> pathJoin(root, os.getenv("HOME"))
set_value = self.REGEX_QUOTE_SHELL_VAR.sub(self.OS_GETENV_TEMPLATE % r"\1", set_value)
else:
val = quote_str(value)
if resolve_env_vars:
# replace env var with os.getenv statement
# example: $HOME -> os.getenv("HOME")
concat_getenv = self.CONCAT_STR + self.OS_GETENV_TEMPLATE % r"\1" + self.CONCAT_STR
set_value = self.REGEX_SHELL_VAR.sub(concat_getenv, set_value)
set_value = self.CONCAT_STR.join([
# quote any substrings that are not a os.getenv Lua statement
x if x.startswith(self.OS_GETENV_TEMPLATE[:10]) else quote_str(x)
for x in set_value.strip(self.CONCAT_STR).split(self.CONCAT_STR)
])

env_setter = 'pushenv' if use_pushenv else 'setenv'

return '%s("%s", %s)\n' % (env_setter, key, val)
return f'{env_setter}("{key}", {set_value})\n'

def swap_module(self, mod_name_out, mod_name_in, guarded=True):
"""
Expand Down
30 changes: 16 additions & 14 deletions test/framework/easyblock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -462,7 +462,7 @@ def test_make_module_req(self):
elif get_module_syntax() == 'Lua':
self.assertTrue(re.search(r'^prepend_path\("CLASSPATH", pathJoin\(root, "bla.jar"\)\)$', guess, re.M))
self.assertTrue(re.search(r'^prepend_path\("CLASSPATH", pathJoin\(root, "foo.jar"\)\)$', guess, re.M))
self.assertTrue(re.search(r'^prepend_path\("MANPATH", pathJoin\(root, "share/man"\)\)$', guess, re.M))
self.assertTrue(re.search(r'^prepend_path\("MANPATH", pathJoin\(root, "share", "man"\)\)$', guess, re.M))
self.assertIn('prepend_path("CMAKE_PREFIX_PATH", root)', guess)
# bin/ is not added to $PATH if it doesn't include files
self.assertFalse(re.search(r'^prepend_path\("PATH", pathJoin\(root, "bin"\)\)$', guess, re.M))
Expand Down Expand Up @@ -573,12 +573,12 @@ def test_make_module_req(self):
r"prepend-path\s+LD_LIBRARY_PATH\s+\$root/lib/pathA\n",
txt, re.M))
elif get_module_syntax() == 'Lua':
self.assertTrue(re.search(r'\nprepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib/pathC"\)\)\n' +
r'prepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib/pathA"\)\)\n' +
r'prepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib/pathB"\)\)\n',
self.assertTrue(re.search(r'\nprepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib", "pathC"\)\)\n' +
r'prepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib", "pathA"\)\)\n' +
r'prepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib", "pathB"\)\)\n',
txt, re.M))
self.assertFalse(re.search(r'\nprepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib/pathB"\)\)\n' +
r'prepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib/pathA"\)\)\n',
self.assertFalse(re.search(r'\nprepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib", "pathB"\)\)\n' +
r'prepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "lib", "pathA"\)\)\n',
txt, re.M))
else:
self.fail("Unknown module syntax: %s" % get_module_syntax())
Expand Down Expand Up @@ -651,10 +651,12 @@ def test_make_module_req(self):
self.assertTrue(re.search(r"^prepend-path\s+LD_LIBRARY_PATH\s+\$root/libraries/intel64_lin$", txt, re.M))
self.assertTrue(re.search(r"^prepend-path\s+LIBRARY_PATH\s+\$root/libraries/intel64_lin\n$", txt, re.M))
elif get_module_syntax() == 'Lua':
self.assertTrue(re.search(r'^prepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "libraries/intel64_lin"\)\)$',
txt, re.M))
self.assertTrue(re.search(r'^prepend_path\("LIBRARY_PATH", pathJoin\(root, "libraries/intel64_lin"\)\)$',
txt, re.M))
self.assertTrue(re.search(
r'^prepend_path\("LD_LIBRARY_PATH", pathJoin\(root, "libraries", "intel64_lin"\)\)$', txt, re.M
))
self.assertTrue(re.search(
r'^prepend_path\("LIBRARY_PATH", pathJoin\(root, "libraries", "intel64_lin"\)\)$', txt, re.M
))
else:
self.fail("Unknown module syntax: %s" % get_module_syntax())

Expand Down Expand Up @@ -712,7 +714,7 @@ def test_make_module_req(self):
expected_patterns = [
r"^append[-_]path.*TEST_VAR_CUSTOM.*root.*foo.*",
r"^prepend[-_]path.*CPATH.*root.*include.*",
r"^prepend[-_]path.*CPATH.*root.*include/foo.*",
r"^prepend[-_]path.*CPATH.*root.*include.*foo.*",
r"^prepend[-_]path.*LD_LIBRARY_PATH.*root.*lib",
r"^prepend[-_]path.*LD_LIBRARY_PATH.*root.*foo",
r"^prepend[-_]path.*TEST_VAR.*root.*foo",
Expand Down Expand Up @@ -868,12 +870,12 @@ def test_make_module_extra(self):
expected_default = re.compile(r'\n'.join([
r'setenv\("EBROOTPI", root\)',
r'setenv\("EBVERSIONPI", "3.14"\)',
r'setenv\("EBDEVELPI", pathJoin\(root, "easybuild/pi-3.14-gompi-2018a-easybuild-devel"\)\)',
r'setenv\("EBDEVELPI", pathJoin\(root, "easybuild", "pi-3.14-gompi-2018a-easybuild-devel"\)\)',
]))
expected_alt = re.compile(r'\n'.join([
r'setenv\("EBROOTPI", "/opt/software/tau/6.28"\)',
r'setenv\("EBVERSIONPI", "6.28"\)',
r'setenv\("EBDEVELPI", pathJoin\(root, "easybuild/pi-3.14-gompi-2018a-easybuild-devel"\)\)',
r'setenv\("EBDEVELPI", pathJoin\(root, "easybuild", "pi-3.14-gompi-2018a-easybuild-devel"\)\)',
]))
else:
self.fail("Unknown module syntax: %s" % get_module_syntax())
Expand Down Expand Up @@ -1614,7 +1616,7 @@ def test_make_module_step(self):
if val == '':
full_val = 'root'
else:
full_val = fr'pathJoin\(root, "{val}"\)'
full_val = fr'pathJoin\(root, "{val.replace("/", ".*")}"\)'
regex = re.compile(fr'^{placement}_path\("{key}", {full_val}{delim_lua}\)$', re.M)
else:
self.fail(f"Unknown module syntax: {get_module_syntax()}")
Expand Down
Loading