From bb4bf27062d58023b02ed1526319e98a1074605e Mon Sep 17 00:00:00 2001
From: Tom Payerle <payerle@umd.edu>
Date: Mon, 19 Apr 2021 13:06:49 -0400
Subject: [PATCH 1/2] Add fix_sys_path stuff

---
 envmodules.py | 56 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/envmodules.py b/envmodules.py
index 40bd008..557a177 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.environ['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.environ['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
+
+

From 6e44095362d12cb5f33119fa38dc4b1131b5e2db Mon Sep 17 00:00:00 2001
From: Tom Payerle <payerle@umd.edu>
Date: Mon, 19 Apr 2021 21:37:01 -0400
Subject: [PATCH 2/2] Fix issues when PYTHONPATH unset

---
 envmodules.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/envmodules.py b/envmodules.py
index 557a177..5cc0431 100644
--- a/envmodules.py
+++ b/envmodules.py
@@ -6,7 +6,7 @@
 
 __version__ = '0.1'
 
-_cached_pythonpath = os.environ['PYTHONPATH']
+_cached_pythonpath = os.getenv('PYTHONPATH', '')
 _auto_fix_sys_path = 0
 
 def _modulecmd(*args):
@@ -40,7 +40,7 @@ def get_auto_fix_sys_path():
 
 def fix_sys_path():
     global _cached_pythonpath
-    pypath = os.environ['PYTHONPATH']
+    pypath = os.getenv('PYTHONPATH', '')
     if pypath == _cached_pythonpath:
         # Nothing to do if PYTHONPATH has not changed
         return