Python: query distutils to find site-packages directory (#24095)

Third-party Python libraries may be installed in one of several directories:

1. `lib/pythonX.Y/site-packages` for Spack-installed Python
2. `lib64/pythonX.Y/site-packages` for system Python on RHEL/CentOS/Fedora
3. `lib/pythonX/dist-packages` for system Python on Debian/Ubuntu

Previously, Spack packages were hard-coded to use the (1). Now, we query the Python installation itself and ask it which to use. Ever since #21446 this is how we've been determining where to install Python libraries anyway.

Note: there are still many packages that are hard-coded to use (1). I can change them in this PR, but I don't have the bandwidth to test all of them.

* Python: handle dist-packages and site-packages
* Query Python to find site-packages directory
* Add try-except statements for when distutils isn't installed
* Catch more errors
* Fix root directory used in import tests
* Rely on site_packages_dir property

4 files changed, 145 insertions, 90 deletions

diff --git a/lib/spack/spack/bootstrap.py b/lib/spack/spack/bootstrap.py
index 96f11b631d..896c1cad4e 100644
--- a/lib/spack/spack/bootstrap.py
+++ b/lib/spack/spack/bootstrap.py
@@ -95,10 +95,8 @@ def make_module_available(module, spec=None, install=False):
         # TODO: make sure run-environment is appropriate
         module_path = os.path.join(ispec.prefix,
                                    ispec['python'].package.site_packages_dir)
-        module_path_64 = module_path.replace('/lib/', '/lib64/')
         try:
             sys.path.append(module_path)
-            sys.path.append(module_path_64)
             __import__(module)
             return
         except ImportError:
@@ -122,10 +120,8 @@ def make_module_available(module, spec=None, install=False):
 
     module_path = os.path.join(spec.prefix,
                                spec['python'].package.site_packages_dir)
-    module_path_64 = module_path.replace('/lib/', '/lib64/')
     try:
         sys.path.append(module_path)
-        sys.path.append(module_path_64)
         __import__(module)
         return
     except ImportError:
diff --git a/lib/spack/spack/build_systems/python.py b/lib/spack/spack/build_systems/python.py
index b4f55ab707..27b1d4a10c 100644
--- a/lib/spack/spack/build_systems/python.py
+++ b/lib/spack/spack/build_systems/python.py
@@ -127,24 +127,22 @@ class PythonPackage(PackageBase):
             list: list of strings of module names
         """
         modules = []
+        root = self.spec['python'].package.get_python_lib(prefix=self.prefix)
 
-        # Python libraries may be installed in lib or lib64
-        # See issues #18520 and #17126
-        for lib in ['lib', 'lib64']:
-            root = os.path.join(self.prefix, lib, 'python{0}'.format(
-                self.spec['python'].version.up_to(2)), 'site-packages')
-            # Some Python libraries are packages: collections of modules
-            # distributed in directories containing __init__.py files
-            for path in find(root, '__init__.py', recursive=True):
-                modules.append(path.replace(root + os.sep, '', 1).replace(
-                    os.sep + '__init__.py', '').replace('/', '.'))
-            # Some Python libraries are modules: individual *.py files
-            # found in the site-packages directory
-            for path in find(root, '*.py', recursive=False):
-                modules.append(path.replace(root + os.sep, '', 1).replace(
-                    '.py', '').replace('/', '.'))
+        # Some Python libraries are packages: collections of modules
+        # distributed in directories containing __init__.py files
+        for path in find(root, '__init__.py', recursive=True):
+            modules.append(path.replace(root + os.sep, '', 1).replace(
+                os.sep + '__init__.py', '').replace('/', '.'))
+
+        # Some Python libraries are modules: individual *.py files
+        # found in the site-packages directory
+        for path in find(root, '*.py', recursive=False):
+            modules.append(path.replace(root + os.sep, '', 1).replace(
+                '.py', '').replace('/', '.'))
 
         tty.debug('Detected the following modules: {0}'.format(modules))
+
         return modules
 
     def setup_file(self):
@@ -254,15 +252,12 @@ class PythonPackage(PackageBase):
         # Get all relative paths since we set the root to `prefix`
         # We query the python with which these will be used for the lib and inc
         # directories. This ensures we use `lib`/`lib64` as expected by python.
-        python = spec['python'].package.command
-        command_start = 'print(distutils.sysconfig.'
-        commands = ';'.join([
-            'import distutils.sysconfig',
-            command_start + 'get_python_lib(plat_specific=False, prefix=""))',
-            command_start + 'get_python_lib(plat_specific=True, prefix=""))',
-            command_start + 'get_python_inc(plat_specific=True, prefix=""))'])
-        pure_site_packages_dir, plat_site_packages_dir, inc_dir = python(
-            '-c', commands, output=str, error=str).strip().split('\n')
+        pure_site_packages_dir = spec['python'].package.get_python_lib(
+            plat_specific=False, prefix='')
+        plat_site_packages_dir = spec['python'].package.get_python_lib(
+            plat_specific=True, prefix='')
+        inc_dir = spec['python'].package.get_python_inc(
+            plat_specific=True, prefix='')
 
         args += ['--root=%s' % prefix,
                  '--install-purelib=%s' % pure_site_packages_dir,
diff --git a/lib/spack/spack/build_systems/sip.py b/lib/spack/spack/build_systems/sip.py
index c32c46ec46..744989e2d6 100644
--- a/lib/spack/spack/build_systems/sip.py
+++ b/lib/spack/spack/build_systems/sip.py
@@ -64,24 +64,22 @@ class SIPPackage(PackageBase):
             list: list of strings of module names
         """
         modules = []
+        root = self.spec['python'].package.get_python_lib(prefix=self.prefix)
 
-        # Python libraries may be installed in lib or lib64
-        # See issues #18520 and #17126
-        for lib in ['lib', 'lib64']:
-            root = os.path.join(self.prefix, lib, 'python{0}'.format(
-                self.spec['python'].version.up_to(2)), 'site-packages')
-            # Some Python libraries are packages: collections of modules
-            # distributed in directories containing __init__.py files
-            for path in find(root, '__init__.py', recursive=True):
-                modules.append(path.replace(root + os.sep, '', 1).replace(
-                    os.sep + '__init__.py', '').replace('/', '.'))
-            # Some Python libraries are modules: individual *.py files
-            # found in the site-packages directory
-            for path in find(root, '*.py', recursive=False):
-                modules.append(path.replace(root + os.sep, '', 1).replace(
-                    '.py', '').replace('/', '.'))
+        # Some Python libraries are packages: collections of modules
+        # distributed in directories containing __init__.py files
+        for path in find(root, '__init__.py', recursive=True):
+            modules.append(path.replace(root + os.sep, '', 1).replace(
+                os.sep + '__init__.py', '').replace('/', '.'))
+
+        # Some Python libraries are modules: individual *.py files
+        # found in the site-packages directory
+        for path in find(root, '*.py', recursive=False):
+            modules.append(path.replace(root + os.sep, '', 1).replace(
+                '.py', '').replace('/', '.'))
 
         tty.debug('Detected the following modules: {0}'.format(modules))
+
         return modules
 
     def python(self, *args, **kwargs):
diff --git a/var/spack/repos/builtin/packages/python/package.py b/var/spack/repos/builtin/packages/python/package.py
index 62805e12a1..61f42aaad3 100644
--- a/var/spack/repos/builtin/packages/python/package.py
+++ b/var/spack/repos/builtin/packages/python/package.py
@@ -711,23 +711,53 @@ class Python(AutotoolsPackage):
 
         return self.command('-c', cmd, output=str).strip()
 
-    def get_python_inc(self):
+    def get_python_inc(self, plat_specific=False, prefix=None):
         """Return the directory for either the general or platform-dependent C
         include files. Wrapper around ``distutils.sysconfig.get_python_inc()``.
+
+        Parameters:
+            plat_specific (bool): if true, the platform-dependent include directory
+                is returned, else the platform-independent directory is returned
+            prefix (str): prefix to use instead of ``distutils.sysconfig.PREFIX``
+
+        Returns:
+            str: include files directory
         """
+        # Wrap strings in quotes
+        if prefix is not None:
+            prefix = '"{0}"'.format(prefix)
+
+        args = 'plat_specific={0}, prefix={1}'.format(plat_specific, prefix)
 
         cmd = 'from distutils.sysconfig import get_python_inc; '
-        cmd += self.print_string('get_python_inc()')
+        cmd += self.print_string('get_python_inc({0})'.format(args))
 
         return self.command('-c', cmd, output=str).strip()
 
-    def get_python_lib(self):
+    def get_python_lib(self, plat_specific=False, standard_lib=False, prefix=None):
         """Return the directory for either the general or platform-dependent
-        library installation. Wrapper around
-        ``distutils.sysconfig.get_python_lib()``."""
+        library installation. Wrapper around ``distutils.sysconfig.get_python_lib()``.
+
+        Parameters:
+            plat_specific (bool): if true, the platform-dependent library directory
+                is returned, else the platform-independent directory is returned
+            standard_lib (bool): if true, the directory for the standard library is
+                returned rather than the directory for the installation of
+                third-party extensions
+            prefix (str): prefix to use instead of ``distutils.sysconfig.PREFIX``
+
+        Returns:
+            str: library installation directory
+        """
+        # Wrap strings in quotes
+        if prefix is not None:
+            prefix = '"{0}"'.format(prefix)
+
+        args = 'plat_specific={0}, standard_lib={1}, prefix={2}'.format(
+            plat_specific, standard_lib, prefix)
 
         cmd = 'from distutils.sysconfig import get_python_lib; '
-        cmd += self.print_string('get_python_lib()')
+        cmd += self.print_string('get_python_lib({0})'.format(args))
 
         return self.command('-c', cmd, output=str).strip()
 
@@ -827,16 +857,71 @@ class Python(AutotoolsPackage):
         return headers
 
     @property
-    def python_lib_dir(self):
-        return join_path('lib', 'python{0}'.format(self.version.up_to(2)))
+    def python_include_dir(self):
+        """Directory for the include files.
+
+        On most systems, and for Spack-installed Python, this will look like:
+
+            ``include/pythonX.Y``
+
+        However, some systems append a ``m`` to the end of this path.
+
+        Returns:
+            str: include files directory
+        """
+        try:
+            return self.get_python_inc(prefix='')
+        except (ProcessError, RuntimeError):
+            return os.path.join('include', 'python{0}'.format(self.version.up_to(2)))
 
     @property
-    def python_include_dir(self):
-        return join_path('include', 'python{0}'.format(self.version.up_to(2)))
+    def python_lib_dir(self):
+        """Directory for the standard library.
+
+        On most systems, and for Spack-installed Python, this will look like:
+
+            ``lib/pythonX.Y``
+
+        On RHEL/CentOS/Fedora, when using the system Python, this will look like:
+
+            ``lib64/pythonX.Y``
+
+        On Debian/Ubuntu, when using the system Python, this will look like:
+
+            ``lib/pythonX``
+
+        Returns:
+            str: standard library directory
+        """
+        try:
+            return self.get_python_lib(standard_lib=True, prefix='')
+        except (ProcessError, RuntimeError):
+            return os.path.join('lib', 'python{0}'.format(self.version.up_to(2)))
 
     @property
     def site_packages_dir(self):
-        return join_path(self.python_lib_dir, 'site-packages')
+        """Directory where third-party extensions should be installed.
+
+        On most systems, and for Spack-installed Python, this will look like:
+
+            ``lib/pythonX.Y/site-packages``
+
+        On RHEL/CentOS/Fedora, when using the system Python, this will look like:
+
+            ``lib64/pythonX.Y/site-packages``
+
+        On Debian/Ubuntu, when using the system Python, this will look like:
+
+            ``lib/pythonX/dist-packages``
+
+        Returns:
+            str: site-packages directory
+        """
+        try:
+            return self.get_python_lib(prefix='')
+        except (ProcessError, RuntimeError):
+            return os.path.join(
+                'lib', 'python{0}'.format(self.version.up_to(2)), 'site-packages')
 
     @property
     def easy_install_file(self):
@@ -848,8 +933,8 @@ class Python(AutotoolsPackage):
 
     def setup_dependent_build_environment(self, env, dependent_spec):
         """Set PYTHONPATH to include the site-packages directory for the
-        extension and any other python extensions it depends on."""
-
+        extension and any other python extensions it depends on.
+        """
         # If we set PYTHONHOME, we must also ensure that the corresponding
         # python is found in the build environment. This to prevent cases
         # where a system provided python is run against the standard libraries
@@ -860,18 +945,10 @@ class Python(AutotoolsPackage):
         if not is_system_path(path):
             env.prepend_path('PATH', path)
 
-        python_paths = []
-        for d in dependent_spec.traverse(deptype=('build', 'run', 'test')):
+        for d in dependent_spec.traverse(deptype=('build', 'run', 'test'), root=True):
             if d.package.extends(self.spec):
-                # Python libraries may be installed in lib or lib64
-                # See issues #18520 and #17126
-                for lib in ['lib', 'lib64']:
-                    python_paths.append(join_path(
-                        d.prefix, lib, 'python' + str(self.version.up_to(2)),
-                        'site-packages'))
-
-        pythonpath = ':'.join(python_paths)
-        env.set('PYTHONPATH', pythonpath)
+                env.prepend_path('PYTHONPATH', join_path(
+                    d.prefix, self.site_packages_dir))
 
         # We need to make sure that the extensions are compiled and linked with
         # the Spack wrapper. Paths to the executables that are used for these
@@ -929,27 +1006,16 @@ class Python(AutotoolsPackage):
                 env.set(link_var, new_link)
 
     def setup_dependent_run_environment(self, env, dependent_spec):
-        python_paths = []
-        for d in dependent_spec.traverse(deptype='run'):
+        """Set PYTHONPATH to include the site-packages directory for the
+        extension and any other python extensions it depends on.
+        """
+        for d in dependent_spec.traverse(deptype=('run'), root=True):
             if d.package.extends(self.spec):
-                # Python libraries may be installed in lib or lib64
-                # See issues #18520 and #17126
-                for lib in ['lib', 'lib64']:
-                    root = join_path(
-                        d.prefix, lib, 'python' + str(self.version.up_to(2)),
-                        'site-packages')
-                    if os.path.exists(root):
-                        python_paths.append(root)
-
-        pythonpath = ':'.join(python_paths)
-        env.prepend_path('PYTHONPATH', pythonpath)
+                env.prepend_path('PYTHONPATH', join_path(
+                    d.prefix, self.site_packages_dir))
 
     def setup_dependent_package(self, module, dependent_spec):
-        """Called before python modules' install() methods.
-
-        In most cases, extensions will only need to have one line::
-
-        setup_py('install', '--prefix={0}'.format(prefix))"""
+        """Called before python modules' install() methods."""
 
         module.python = self.command
         module.setup_py = Executable(
@@ -978,17 +1044,17 @@ class Python(AutotoolsPackage):
         ignore_arg = args.get('ignore', lambda f: False)
 
         # Always ignore easy-install.pth, as it needs to be merged.
-        patterns = [r'site-packages/easy-install\.pth$']
+        patterns = [r'(site|dist)-packages/easy-install\.pth$']
 
         # Ignore pieces of setuptools installed by other packages.
         # Must include directory name or it will remove all site*.py files.
         if ext_pkg.name != 'py-setuptools':
             patterns.extend([
                 r'bin/easy_install[^/]*$',
-                r'site-packages/setuptools[^/]*\.egg$',
-                r'site-packages/setuptools\.pth$',
-                r'site-packages/site[^/]*\.pyc?$',
-                r'site-packages/__pycache__/site[^/]*\.pyc?$'
+                r'(site|dist)-packages/setuptools[^/]*\.egg$',
+                r'(site|dist)-packages/setuptools\.pth$',
+                r'(site|dist)-packages/site[^/]*\.pyc?$',
+                r'(site|dist)-packages/__pycache__/site[^/]*\.pyc?$'
             ])
         if ext_pkg.name != 'py-pygments':
             patterns.append(r'bin/pygmentize$')