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

1 files changed, 113 insertions, 47 deletions

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$')