diff --git a/envmodules.py b/envmodules.py index 40bd008..5cc0431 100644 --- a/envmodules.py +++ b/envmodules.py @@ -1,10 +1,14 @@ """Python wrapper around environment modules ("module load") """ import os +import sys from subprocess import run, PIPE __version__ = '0.1' +_cached_pythonpath = os.getenv('PYTHONPATH', '') +_auto_fix_sys_path = 0 + def _modulecmd(*args): cmd = ['modulecmd', 'python'] + list(args) res = run(cmd, stdout=PIPE, stderr=PIPE) @@ -18,8 +22,60 @@ def _modulecmd(*args): if code: exec(code, {'os': os}) + if _auto_fix_sys_path: + fix_sys_path() + def load(*args): _modulecmd('load', *args) def unload(*args): _modulecmd('unload', *args) + +def set_auto_fix_sys_path(new): + global _auto_fix_sys_path + _auto_fix_sys_path = bool(new) + +def get_auto_fix_sys_path(): + return _auto_fix_sys_path + +def fix_sys_path(): + global _cached_pythonpath + pypath = os.getenv('PYTHONPATH', '') + if pypath == _cached_pythonpath: + # Nothing to do if PYTHONPATH has not changed + return + + # PYTHONPATH has changed, see what changed + oldpypaths = _cached_pythonpath.split(':') + newpypaths = pypath.split(':') + + oldpyset = set(oldpypaths) + newpyset = set(newpypaths) + + oldpyonly = [] + newpyonly = [] + + # Find elements only in oldpypaths + for path in oldpypaths: + if path not in newpyset: + oldpyonly.append(path) + + # Find elements only in newpypaths + # We reverse the order for this list (prepend instead of append) + for path in newpypaths: + if path not in oldpyset: + newpyonly.insert(0, path) + + # Modify sys.path + # Remove from sys.path paths in oldpypaths only + for path in oldpyonly: + sys.path.remove(path) + + # Prepend to sys.path paths in newpypaths only + for path in newpyonly: + sys.path.insert(0, path) + + # Cache the new PYTHONPATH + _cached_pythonpath = pypath + +