Merge pull request #4539 from Micket/ebpythonprefixes2

prepend to `$PYTHONPATH` or `$EBPYTHONPREFIXES` in generated module files by automatically scanning for python site package directories

Author: lexming

easybuild/framework/easyblock.py

@@ -72,7 +72,7 @@
 from easybuild.tools.build_details import get_build_stats
 from easybuild.tools.build_log import EasyBuildError, EasyBuildExit, dry_run_msg, dry_run_warning, dry_run_set_dirs
 from easybuild.tools.build_log import print_error, print_msg, print_warning
-from easybuild.tools.config import CHECKSUM_PRIORITY_JSON, DEFAULT_ENVVAR_USERS_MODULES
+from easybuild.tools.config import CHECKSUM_PRIORITY_JSON, DEFAULT_ENVVAR_USERS_MODULES, PYTHONPATH, EBPYTHONPREFIXES
 from easybuild.tools.config import FORCE_DOWNLOAD_ALL, FORCE_DOWNLOAD_PATCHES, FORCE_DOWNLOAD_SOURCES
 from easybuild.tools.config import EASYBUILD_SOURCES_URL # noqa
 from easybuild.tools.config import build_option, build_path, get_log_filename, get_repository, get_repositorypath
@@ -1390,6 +1390,49 @@ def make_module_description(self):
         """
         return self.module_generator.get_description()
 
+    def make_module_pythonpath(self):
+        """
+        Add lines for module file to update $PYTHONPATH or $EBPYTHONPREFIXES,
+        if they aren't already present and the standard lib/python*/site-packages subdirectory exists
+        """
+        lines = []
+        if not os.path.isfile(os.path.join(self.installdir, 'bin', 'python')):  # only needed when not a python install
+            python_subdir_pattern = os.path.join(self.installdir, 'lib', 'python*', 'site-packages')
+            candidate_paths = (os.path.relpath(path, self.installdir) for path in glob.glob(python_subdir_pattern))
+            python_paths = [path for path in candidate_paths if re.match(r'lib/python\d+\.\d+/site-packages', path)]
+
+            # determine whether Python is a runtime dependency;
+            # if so, we assume it was installed with EasyBuild, and hence is aware of $EBPYTHONPREFIXES
+            runtime_deps = [dep['name'] for dep in self.cfg.dependencies(runtime_only=True)]
+
+            # don't use $EBPYTHONPREFIXES unless we can and it's preferred or necesary (due to use of multi_deps)
+            use_ebpythonprefixes = False
+            multi_deps = self.cfg['multi_deps']
+
+            if 'Python' in runtime_deps:
+                self.log.info("Found Python runtime dependency, so considering $EBPYTHONPREFIXES...")
+
+                if build_option('prefer_python_search_path') == EBPYTHONPREFIXES:
+                    self.log.info("Preferred Python search path is $EBPYTHONPREFIXES, so using that")
+                    use_ebpythonprefixes = True
+
+            elif multi_deps and 'Python' in multi_deps:
+                self.log.info("Python is listed in 'multi_deps', so using $EBPYTHONPREFIXES instead of $PYTHONPATH")
+                use_ebpythonprefixes = True
+
+            if python_paths:
+                # add paths unless they were already added
+                if use_ebpythonprefixes:
+                    path = ''  # EBPYTHONPREFIXES are relative to the install dir
+                    if path not in self.module_generator.added_paths_per_key[EBPYTHONPREFIXES]:
+                        lines.append(self.module_generator.prepend_paths(EBPYTHONPREFIXES, path))
+                else:
+                    for python_path in python_paths:
+                        if python_path not in self.module_generator.added_paths_per_key[PYTHONPATH]:
+                            lines.append(self.module_generator.prepend_paths(PYTHONPATH, python_path))
+
+        return lines
+
     def make_module_extra(self, altroot=None, altversion=None):
         """
         Set extra stuff in module file, e.g. $EBROOT*, $EBVERSION*, etc.
@@ -1438,6 +1481,9 @@ def make_module_extra(self, altroot=None, altversion=None):
                                      value, type(value))
             lines.append(self.module_generator.append_paths(key, value, allow_abs=self.cfg['allow_append_abs_path']))
 
+        # add lines to update $PYTHONPATH or $EBPYTHONPREFIXES
+        lines.extend(self.make_module_pythonpath())
+
         modloadmsg = self.cfg['modloadmsg']
         if modloadmsg:
             # add trailing newline to prevent that shell prompt is 'glued' to module load message

easybuild/framework/easyconfig/easyconfig.py

@@ -1105,15 +1105,19 @@ def filter_deps(self, deps):
 
         return retained_deps
 
-    def dependencies(self, build_only=False):
+    def dependencies(self, build_only=False, runtime_only=False):
         """
         Returns an array of parsed dependencies (after filtering, if requested)
         dependency = {'name': '', 'version': '', 'system': (False|True), 'versionsuffix': '', 'toolchain': ''}
         Iterable builddependencies are flattened when not iterating.
 
         :param build_only: only return build dependencies, discard others
+        :param runtime_only: only return runtime dependencies, discard others
         """
-        deps = self.builddependencies()
+        if runtime_only:
+            deps = []
+        else:
+            deps = self.builddependencies()
 
         if not build_only:
             # use += rather than .extend to get a new list rather than updating list of build deps in place...

easybuild/tools/config.py

@@ -175,6 +175,10 @@
 OUTPUT_STYLE_RICH = 'rich'
 OUTPUT_STYLES = (OUTPUT_STYLE_AUTO, OUTPUT_STYLE_BASIC, OUTPUT_STYLE_NO_COLOR, OUTPUT_STYLE_RICH)
 
+PYTHONPATH = 'PYTHONPATH'
+EBPYTHONPREFIXES = 'EBPYTHONPREFIXES'
+PYTHON_SEARCH_PATH_TYPES = [PYTHONPATH, EBPYTHONPREFIXES]
+
 
 class Singleton(ABCMeta):
     """Serves as metaclass for classes that should implement the Singleton pattern.
@@ -407,6 +411,9 @@ def mk_full_default_path(name, prefix=DEFAULT_PREFIX):
     OUTPUT_STYLE_AUTO: [
         'output_style',
     ],
+    PYTHONPATH: [
+        'prefer_python_search_path',
+    ]
 }
 # build option that do not have a perfectly matching command line option
 BUILD_OPTIONS_OTHER = {

easybuild/tools/module_generator.py

@@ -39,6 +39,7 @@
 import os
 import re
 import tempfile
+from collections import defaultdict
 from contextlib import contextmanager
 from easybuild.tools import LooseVersion
 from textwrap import wrap
@@ -153,7 +154,7 @@ def start_module_creation(self):
             raise EasyBuildError('Module creation already in process. '
                                  'You cannot create multiple modules at the same time!')
         # Mapping of keys/env vars to paths already added
-        self.added_paths_per_key = dict()
+        self.added_paths_per_key = defaultdict(set)
         txt = self.MODULE_SHEBANG
         if txt:
             txt += '\n'
@@ -212,7 +213,7 @@ def _filter_paths(self, key, paths):
             print_warning('Module creation has not been started. Call start_module_creation first!')
             return paths
 
-        added_paths = self.added_paths_per_key.setdefault(key, set())
+        added_paths = self.added_paths_per_key[key]
         # paths can be a string
         if isinstance(paths, str):
             if paths in added_paths:

easybuild/tools/options.py

@@ -78,6 +78,7 @@
 from easybuild.tools.config import OUTPUT_STYLE_AUTO, OUTPUT_STYLES, WARN
 from easybuild.tools.config import get_pretend_installpath, init, init_build_options, mk_full_default_path
 from easybuild.tools.config import BuildOptions, ConfigurationVariables
+from easybuild.tools.config import PYTHON_SEARCH_PATH_TYPES, PYTHONPATH
 from easybuild.tools.configobj import ConfigObj, ConfigObjError
 from easybuild.tools.docs import FORMAT_JSON, FORMAT_MD, FORMAT_RST, FORMAT_TXT
 from easybuild.tools.docs import avail_cfgfile_constants, avail_easyconfig_constants, avail_easyconfig_licenses
@@ -490,6 +491,10 @@ def override_options(self):
                                             None, 'store_true', False),
             'pre-create-installdir': ("Create installation directory before submitting build jobs",
                                       None, 'store_true', True),
+            'prefer-python-search-path': (("Prefer using specified environment variable when possible to specify where"
+                                           " Python packages were installed; see also "
+                                           "https://docs.easybuild.io/python-search-path"),
+                                          'choice', 'store_or_None', PYTHONPATH, PYTHON_SEARCH_PATH_TYPES),
             'pretend': (("Does the build/installation in a test directory located in $HOME/easybuildinstall"),
                         None, 'store_true', False, 'p'),
             'read-only-installdir': ("Set read-only permissions on installation directory after installation",

test/framework/toy_build.py

@@ -4238,6 +4238,70 @@ def test_eb_error(self):
             stderr = stderr.getvalue()
             self.assertTrue(regex.search(stderr), f"Pattern '{regex.pattern}' should be found in {stderr}")
 
+    def test_toy_python(self):
+        """
+        Test whether $PYTHONPATH or $EBPYTHONPREFIXES are set correctly.
+        """
+        # generate fake Python modules that we can use as runtime dependency for toy
+        # (required condition for use of $EBPYTHONPREFIXES)
+        fake_mods_path = os.path.join(self.test_prefix, 'modules')
+        for pyver in ('2.7', '3.6'):
+            fake_python_mod = os.path.join(fake_mods_path, 'Python', pyver)
+            if get_module_syntax() == 'Lua':
+                fake_python_mod += '.lua'
+                write_file(fake_python_mod, '')
+            else:
+                write_file(fake_python_mod, '#%Module')
+        self.modtool.use(fake_mods_path)
+
+        test_ecs = os.path.join(os.path.dirname(__file__), 'easyconfigs', 'test_ecs')
+        toy_ec = os.path.join(test_ecs, 't', 'toy', 'toy-0.0.eb')
+
+        test_ec_txt = read_file(toy_ec)
+        test_ec_txt += "\npostinstallcmds.append('mkdir -p %(installdir)s/lib/python3.6/site-packages')"
+        test_ec_txt += "\npostinstallcmds.append('touch %(installdir)s/lib/python3.6/site-packages/foo.py')"
+
+        test_ec = os.path.join(self.test_prefix, 'test.eb')
+        write_file(test_ec, test_ec_txt)
+        self.run_test_toy_build_with_output(ec_file=test_ec)
+
+        toy_mod = os.path.join(self.test_installpath, 'modules', 'all', 'toy', '0.0')
+        if get_module_syntax() == 'Lua':
+            toy_mod += '.lua'
+        toy_mod_txt = read_file(toy_mod)
+
+        pythonpath_regex = re.compile('^prepend.path.*PYTHONPATH.*lib/python3.6/site-packages', re.M)
+
+        self.assertTrue(pythonpath_regex.search(toy_mod_txt),
+                        f"Pattern '{pythonpath_regex.pattern}' found in: {toy_mod_txt}")
+
+        # also check when opting in to use $EBPYTHONPREFIXES instead of $PYTHONPATH
+        args = ['--prefer-python-search-path=EBPYTHONPREFIXES']
+        self.run_test_toy_build_with_output(ec_file=test_ec, extra_args=args)
+        toy_mod_txt = read_file(toy_mod)
+        # if Python is not listed as a runtime dependency then $PYTHONPATH is still used,
+        # because the Python dependency used must be aware of $EBPYTHONPREFIXES
+        # (see sitecustomize.py installed by Python easyblock)
+        self.assertTrue(pythonpath_regex.search(toy_mod_txt),
+                        f"Pattern '{pythonpath_regex.pattern}' found in: {toy_mod_txt}")
+
+        # if Python is listed as runtime dependency, then $EBPYTHONPREFIXES is used if it's preferred
+        write_file(test_ec, test_ec_txt + "\ndependencies = [('Python', '3.6', '', SYSTEM)]")
+        self.run_test_toy_build_with_output(ec_file=test_ec, extra_args=args)
+        toy_mod_txt = read_file(toy_mod)
+
+        ebpythonprefixes_regex = re.compile('^prepend.path.*EBPYTHONPREFIXES.*root', re.M)
+        self.assertTrue(ebpythonprefixes_regex.search(toy_mod_txt),
+                        f"Pattern '{ebpythonprefixes_regex.pattern}' found in: {toy_mod_txt}")
+
+        # if Python is listed in multi_deps, then $EBPYTHONPREFIXES is used, even if it's not explicitely preferred
+        write_file(test_ec, test_ec_txt + "\nmulti_deps = {'Python': ['2.7', '3.6']}")
+        self.run_test_toy_build_with_output(ec_file=test_ec)
+        toy_mod_txt = read_file(toy_mod)
+
+        self.assertTrue(ebpythonprefixes_regex.search(toy_mod_txt),
+                        f"Pattern '{ebpythonprefixes_regex.pattern}' found in: {toy_mod_txt}")
+
 
 def suite():
     """ return all the tests in this file """