PoC, use dedicated contextes for environment

Author: Crivella

easybuild/main.py

@@ -49,6 +49,9 @@
 
 # IMPORTANT this has to be the first easybuild import as it customises the logging
 #  expect missing log output when this not the case!
+from easybuild import os_hook  # Imported to inject hook that replaces system os with our wrapped version
+os_hook.install_os_hook()
+
 from easybuild.tools.build_log import EasyBuildError, print_error_and_exit, print_msg, print_warning, stop_logging
 from easybuild.tools.build_log import EasyBuildExit
 

easybuild/os_hook.py

@@ -0,0 +1,53 @@
+import importlib
+import importlib.abc
+import importlib.util
+import sys
+import types
+
+
+class OSProxy(types.ModuleType):
+    """Proxy module to intercept os attribute access."""
+    overrides = {}
+    def __init__(self, real_os):
+        super().__init__("os")
+        self._real_os = real_os
+
+    def __getattr__(self, name):
+        # Intercept specific attributes
+        return OSProxy.overrides.get(name, getattr(self._real_os, name))
+
+    def __dir__(self):
+        return dir(self._real_os)
+
+    @classmethod
+    def register_override(cls, name, value):
+        cls.overrides[name] = value
+
+class OSFinder(importlib.abc.MetaPathFinder):
+    """Meta path finder to intercept imports of 'os' and return our proxy."""
+    def find_spec(self, fullname, path, target=None):
+        if fullname == "os":
+            return importlib.util.spec_from_loader(fullname, OSLoader())
+        return None
+
+
+class OSLoader(importlib.abc.Loader):
+    """Loader to create our OSProxy instead of the real os module."""
+    def create_module(self, spec):
+        # Import real os safely
+        sys.meta_path = [f for f in sys.meta_path if not isinstance(f, OSFinder)]
+        real_os = importlib.import_module("os")
+        sys.meta_path.insert(0, OSFinder())
+
+        # Return proxy instead of real module
+        return OSProxy(real_os)
+
+
+def install_os_hook():
+    """Install the os hooking mechanism to intercept imports of 'os' and return our proxy."""
+    sys.meta_path.insert(0, OSFinder())
+
+    # If already imported, replace in place
+    if "os" in sys.modules:
+        real_os = sys.modules["os"]
+        sys.modules["os"] = OSProxy(real_os)

easybuild/tools/environment.py

@@ -32,7 +32,9 @@
 """
 import copy
 import os
+from contextlib import contextmanager
 
+from easybuild.os_hook import OSProxy
 from easybuild.base import fancylogger
 from easybuild.tools.build_log import EasyBuildError, dry_run_msg
 from easybuild.tools.config import build_option
@@ -45,16 +47,35 @@
 
 _log = fancylogger.getLogger('environment', fname=False)
 
-_changes = {}
+_contextes = {'': {}}
+_curr_context = ''
 
+def set_context(context_name, context = {}):
+    """
+    Set context for tracking environment changes.
+    """
+    global _curr_context
+    _curr_context = context_name
+    if context_name not in _contextes:
+        if context is not None:
+            context = copy.deepcopy(context)
+        else:
+            context = {}
+        _contextes[context_name] = context
+
+def get_context():
+    """
+    Return current context for tracking environment changes.
+    """
+    return _contextes[_curr_context]
 
 def write_changes(filename):
     """
     Write current changes to filename and reset environment afterwards
     """
     try:
         with open(filename, 'w') as script:
-            for key, changed_value in _changes.items():
+            for key, changed_value in get_context().items():
                 script.write('export %s=%s\n' % (key, shell_quote(changed_value)))
     except IOError as err:
         raise EasyBuildError("Failed to write to %s: %s", filename, err)
@@ -65,32 +86,79 @@ def reset_changes():
     """
     Reset the changes tracked by this module
     """
-    global _changes
-    _changes = {}
+    get_context().clear()
 
 
 def get_changes():
     """
     Return tracked changes made in environment.
     """
-    return _changes
+    return get_context()
+
+def apply_context(context = None):
+    """Return the current environment with the changes tracked in the context applied.
+
+    Args:
+        context (str, optional): The context to apply. Defaults to the current onee.
+    """
+    if context is None:
+        context = _curr_context
+    changes = get_context()
+    # print(f'Applying context {context} with changes: {changes}')
+    curr_env = ORIG_OS_ENVIRON.copy()
+    for key, changed_value in changes.items():
+        if changed_value is None:
+            curr_env.pop(key, None)
+        else:
+            curr_env[key] = changed_value
+    return curr_env
+
+
+def getvar(key, default=''):
+    """
+    Return value of key in the environment, or default if not found
+    """
+    return get_context().get(key, os._real_os.environ.get(key, default))
+
+
+@contextmanager
+def with_environment(copy_current=False):
+    """Context manager to run code in a dedicated context"""
+    # Get a key that does not exist in _contextes
+    base = '_context_'
+    cnt = 0
+    while (context := f"{base}{cnt}") in _contextes:
+        cnt += 1
+
+    prev_context = _curr_context
+    kwargs = {}
+    if copy_current:
+        kwargs['context'] = get_context()
+    set_context(context, **kwargs)
+    try:
+        yield
+    finally:
+        set_context(prev_context)
+        _contextes.pop(context, None)
 
 
-def setvar(key, value, verbose=True, log_changes=True):
+def setvar(key, value, verbose=True, log_changes=True, force_env=False):
     """
     put key in the environment with value
     tracks added keys until write_changes has been called
 
     :param verbose: include message in dry run output for defining this environment variable
     :param log_changes: show the change in the log
+    :param force_env: if True, also set the variable in os.environ, eg for calls to python functions that rely
+                      on some specific environment such as TMPDIR for tempfile.
     """
     try:
         oldval_info = "previous value: '%s'" % os.environ[key]
     except KeyError:
         oldval_info = "previously undefined"
-    # os.putenv() is not necessary. os.environ will call this.
-    os.environ[key] = value
-    _changes[key] = value
+    if force_env:
+        os._real_os.environ[key] = value
+    get_context()[key] = value
     if log_changes:
         _log.info("Environment variable %s set to %s (%s)", key, value, oldval_info)
 
@@ -101,7 +169,27 @@ def setvar(key, value, verbose=True, log_changes=True):
         dry_run_msg("  export %s=%s" % (key, quoted_value), silent=build_option('silent'))
 
 
-def unset_env_vars(keys, verbose=True):
+def appendvar(key, value, sep=os.pathsep, verbose=True, log_changes=True, force_env=False):
+    """append value to key in the environment, using sep as separator"""
+    oldval = apply_context().get(key)
+    if oldval:
+        newval = oldval + sep + value
+    else:
+        newval = value
+    setvar(key, newval, verbose=verbose, log_changes=log_changes, force_env=force_env)
+
+
+def prependvar(key, value, sep=os.pathsep, verbose=True, log_changes=True, force_env=False):
+    """prepend value to key in the environment, using sep as separator"""
+    oldval = apply_context().get(key)
+    if oldval:
+        newval = value + sep + oldval
+    else:
+        newval = value
+    setvar(key, newval, verbose=verbose, log_changes=log_changes, force_env=force_env)
+
+
+def unset_env_vars(keys, verbose=True, force_env=False):
     """
     Unset the keys given in the environment
     Returns a dict with the old values of the unset keys
@@ -115,7 +203,9 @@ def unset_env_vars(keys, verbose=True):
         if key in os.environ:
             _log.info("Unsetting environment variable %s (value: %s)" % (key, os.environ[key]))
             old_environ[key] = os.environ[key]
-            del os.environ[key]
+            if force_env:
+                del os._real_os.environ[key]
+            get_context()[key] = None
             if verbose and build_option('extended_dry_run'):
                 dry_run_msg("  unset %s  # value was: %s" % (key, old_environ[key]), silent=build_option('silent'))
 
@@ -223,3 +313,29 @@ 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)
+
+
+class MockEnviron(dict):
+    """Hook into os.environ and replace it with calls from this module to track changes to the environment."""
+    def __getitem__(self, key):
+        return getvar(key)
+
+    def __setitem__(self, key, value):
+        setvar(key, value, verbose=False, log_changes=False)
+
+    def __delitem__(self, key):
+        unset_env_vars([key], verbose=False)
+
+    def get(self, key, default=None):
+        return getvar(key, default)
+
+    def copy(self):
+        return apply_context()
+
+    def __deepcopy__(self, memo):
+        return apply_context()
+
+
+OSProxy.register_override('environ', MockEnviron())
+OSProxy.register_override('getenv', lambda key, default=None: getvar(key, default))
+OSProxy.register_override('unsetenv', lambda key: unset_env_vars([key], verbose=False))

easybuild/tools/options.py

@@ -2148,7 +2148,7 @@ def set_tmpdir(tmpdir=None, raise_error=False):
     _log.info("Temporary directory used in this EasyBuild run: %s" % current_tmpdir)
 
     for var in ['TMPDIR', 'TEMP', 'TMP']:
-        env.setvar(var, current_tmpdir, verbose=False)
+        env.setvar(var, current_tmpdir, verbose=False, force_env=True)
 
     # reset to make sure tempfile picks up new temporary directory to use
     tempfile.tempdir = None

easybuild/tools/run.py

@@ -68,6 +68,7 @@
 from easybuild.tools.hooks import RUN_SHELL_CMD, load_hooks, run_hook
 from easybuild.tools.output import COLOR_RED, COLOR_YELLOW, colorize, escape_for_rich, print_error
 from easybuild.tools.utilities import trace_msg
+from easybuild.tools.environment import apply_context
 
 
 _log = fancylogger.getLogger('run', fname=False)
@@ -406,6 +407,9 @@ def to_cmd_str(cmd):
 
         return cmd_str
 
+    if env is None:
+        env = apply_context()
+
     # make sure that qa_patterns is a list of 2-tuples (not a dict, or something else)
     if qa_patterns:
         if not isinstance(qa_patterns, list) or any(not isinstance(x, tuple) or len(x) != 2 for x in qa_patterns):

easybuild/tools/toolchain/toolchain.py

@@ -61,7 +61,7 @@
 from easybuild.base import fancylogger
 from easybuild.tools.build_log import EasyBuildError, dry_run_msg, print_warning
 from easybuild.tools.config import build_option, install_path
-from easybuild.tools.environment import setvar
+from easybuild.tools.environment import setvar, prependvar
 from easybuild.tools.filetools import adjust_permissions, copy_file, find_eb_script, mkdir, read_file, which, write_file
 from easybuild.tools.module_generator import dependencies_for
 from easybuild.tools.modules import get_software_root, get_software_root_env_var_name
@@ -830,7 +830,7 @@ def symlink_commands(self, paths):
         symlink_dir = tempfile.mkdtemp()
 
         # prepend location to symlinks to $PATH
-        setvar('PATH', '%s:%s' % (symlink_dir, os.getenv('PATH')))
+        prependvar('PATH', symlink_dir)
 
         for (path, cmds) in paths.values():
             for cmd in cmds:
@@ -1145,7 +1145,7 @@ def prepare_rpath_wrappers(self, rpath_filter_dirs=None, rpath_include_dirs=None
                 adjust_permissions(cmd_wrapper, stat.S_IXUSR)
 
                 # prepend location to this wrapper to $PATH
-                setvar('PATH', '%s:%s' % (wrapper_dir, os.getenv('PATH')))
+                prependvar('PATH', wrapper_dir)
 
                 self.log.info("RPATH wrapper script for %s: %s (log: %s)", orig_cmd, which(cmd), rpath_wrapper_log)
             else: