WP-hpc-tools/lmod/ipython/startup/lmod_python_fix.py

# -*- python -*-
#from __future__ import print_function
#from subprocess import PIPE, Popen
import os, sys 

sys.path.append('/shared/apps/common/lmod/current/init')

from env_modules_python import module as lmod_module

def module(command, *arguments):
    lmod_module(command, *arguments)
    # BAD: This will cause alteration of PYTHONPATH in a way that may not be desirable.
    sys_path_orig = sys.path
    PYTHONPATH = os.environ.get('PYTHONPATH')
    sys_path_new = PYTHONPATH.split(':') if PYTHONPATH is not None else []
    for p in sys_path_orig:
        if p not in sys_path_new:
            sys_path_new.append(p)
    print(sys_path_new)
    sys.path = sys_path_new

    manual_ld_library_dir = os.environ.get('LMOD_MANUAL_LD_LIBRARY_PATH')
    if manual_ld_library_dir is None:
        return

    for search_path in os.getenv('LD_LIBRARY_PATH').split(':')[::-1]:
        if not os.path.isdir(search_path):
            continue
        if search_path == manual_ld_library_dir:
            continue

        for f in os.listdir(search_path):
            if '.so' in f:
                src = f'{search_path}/{f}'
                tgt = f'{manual_ld_library_dir}/{f}'
                if os.path.islink(tgt):
                    os.unlink(tgt)

                os.symlink(src, tgt)