Bugfix: External Python Extensions (#34202)

Normally when using external packages in concretization, Spack ignores
all dependencies of the external. #33777 updated this logic to attach
a Python Spec to external Python extensions (most py-* packages), but
as implemented there were a couple issues:

* this did not account for concretization groups and could generate
  multiple different python specs for a single DAG
* in some cases this created a fake Python spec with insufficient
  details to be usable (concretization/installation of the
  extension would fail)

This PR addresses both of these issues: 

* For environment specs that are concretized together, external python
  extensions in those specs will all be assigned the same Python spec
* If Spack needs to "invent" a Python spec, then it will have all the
  needed details (e.g. compiler/architecture)

4 files changed, 207 insertions, 20 deletions

diff --git a/lib/spack/spack/build_systems/python.py b/lib/spack/spack/build_systems/python.py
index 3c849cf8fb..d93807ff20 100644
--- a/lib/spack/spack/build_systems/python.py
+++ b/lib/spack/spack/build_systems/python.py
@@ -8,14 +8,19 @@ import re
 import shutil
 from typing import Optional
 
+import archspec
+
 import llnl.util.filesystem as fs
 import llnl.util.lang as lang
 import llnl.util.tty as tty
 
 import spack.builder
+import spack.config
+import spack.detection
 import spack.multimethod
 import spack.package_base
 import spack.spec
+import spack.store
 from spack.directives import build_system, depends_on, extends
 from spack.error import NoHeadersError, NoLibrariesError, SpecError
 from spack.version import Version
@@ -219,7 +224,7 @@ class PythonPackage(PythonExtension):
             name = cls.pypi.split("/")[0]
             return "https://pypi.org/simple/" + name + "/"
 
-    def update_external_dependencies(self):
+    def update_external_dependencies(self, extendee_spec=None):
         """
         Ensure all external python packages have a python dependency
 
@@ -230,16 +235,81 @@ class PythonPackage(PythonExtension):
         """
         # TODO: Include this in the solve, rather than instantiating post-concretization
         if "python" not in self.spec:
-            if "python" in self.spec.root:
+            if extendee_spec:
+                python = extendee_spec
+            elif "python" in self.spec.root:
                 python = self.spec.root["python"]
             else:
-                python = spack.spec.Spec("python")
-                repo = spack.repo.path.repo_for_pkg(python)
-                python.namespace = repo.namespace
-                python._mark_concrete()
-                python.external_path = self.prefix
+                python = self.get_external_python_for_prefix()
+                if not python.concrete:
+                    repo = spack.repo.path.repo_for_pkg(python)
+                    python.namespace = repo.namespace
+
+                    # Ensure architecture information is present
+                    if not python.architecture:
+                        host_platform = spack.platforms.host()
+                        host_os = host_platform.operating_system("default_os")
+                        host_target = host_platform.target("default_target")
+                        python.architecture = spack.spec.ArchSpec(
+                            (str(host_platform), str(host_os), str(host_target))
+                        )
+                    else:
+                        if not python.architecture.platform:
+                            python.architecture.platform = spack.platforms.host()
+                        if not python.architecture.os:
+                            python.architecture.os = "default_os"
+                        if not python.architecture.target:
+                            python.architecture.target = archspec.cpu.host().family.name
+
+                    # Ensure compiler information is present
+                    if not python.compiler:
+                        python.compiler = self.spec.compiler
+
+                    python.external_path = self.spec.external_path
+                    python._mark_concrete()
             self.spec.add_dependency_edge(python, ("build", "link", "run"))
 
+    def get_external_python_for_prefix(self):
+        """
+        For an external package that extends python, find the most likely spec for the python
+        it depends on.
+
+        First search: an "installed" external that shares a prefix with this package
+        Second search: a configured external that shares a prefix with this package
+        Third search: search this prefix for a python package
+
+        Returns:
+          spack.spec.Spec: The external Spec for python most likely to be compatible with self.spec
+        """
+        python_externals_installed = [
+            s for s in spack.store.db.query("python") if s.prefix == self.spec.external_path
+        ]
+        if python_externals_installed:
+            return python_externals_installed[0]
+
+        python_external_config = spack.config.get("packages:python:externals", [])
+        python_externals_configured = [
+            spack.spec.Spec(item["spec"])
+            for item in python_external_config
+            if item["prefix"] == self.spec.external_path
+        ]
+        if python_externals_configured:
+            return python_externals_configured[0]
+
+        python_externals_detection = spack.detection.by_executable(
+            [spack.repo.path.get_pkg_class("python")], path_hints=[self.spec.external_path]
+        )
+
+        python_externals_detected = [
+            d.spec
+            for d in python_externals_detection.get("python", [])
+            if d.prefix == self.spec.external_path
+        ]
+        if python_externals_detected:
+            return python_externals_detected[0]
+
+        raise StopIteration("No external python could be detected for %s to depend on" % self.spec)
+
     @property
     def headers(self):
         """Discover header files in platlib."""
diff --git a/lib/spack/spack/package_base.py b/lib/spack/spack/package_base.py
index 686d490b8b..7c6a2d301b 100644
--- a/lib/spack/spack/package_base.py
+++ b/lib/spack/spack/package_base.py
@@ -914,7 +914,7 @@ class PackageBase(WindowsRPath, PackageViewMixin, metaclass=PackageMeta):
         """
         return self._implement_all_urls_for_version(version)[0]
 
-    def update_external_dependencies(self):
+    def update_external_dependencies(self, extendee_spec=None):
         """
         Method to override in package classes to handle external dependencies
         """
diff --git a/lib/spack/spack/solver/asp.py b/lib/spack/spack/solver/asp.py
index dca4296f56..c6a8ea0c89 100644
--- a/lib/spack/spack/solver/asp.py
+++ b/lib/spack/spack/solver/asp.py
@@ -2245,6 +2245,12 @@ class SpecBuilder(object):
         )
         self._specs[pkg].extra_attributes = spec_info.get("extra_attributes", {})
 
+        # If this is an extension, update the dependencies to include the extendee
+        package = self._specs[pkg].package_class(self._specs[pkg])
+        extendee_spec = package.extendee_spec
+        if extendee_spec:
+            package.update_external_dependencies(self._specs.get(extendee_spec.name, None))
+
     def depends_on(self, pkg, dep, type):
         dependencies = self._specs[pkg].edges_to_dependencies(name=dep)
 
@@ -2311,17 +2317,28 @@ class SpecBuilder(object):
 
     @staticmethod
     def sort_fn(function_tuple):
+        """Ensure attributes are evaluated in the correct order.
+
+        hash attributes are handled first, since they imply entire concrete specs
+        node attributes are handled next, since they instantiate nodes
+        node_compiler attributes are handled next to ensure they come before node_compiler_version
+        external_spec_selected attributes are handled last, so that external extensions can find
+        the concrete specs on which they depend because all nodes are fully constructed before we
+        consider which ones are external.
+        """
         name = function_tuple[0]
         if name == "hash":
-            return (-4, 0)
+            return (-5, 0)
         elif name == "node":
-            return (-3, 0)
+            return (-4, 0)
         elif name == "node_compiler":
-            return (-2, 0)
+            return (-3, 0)
         elif name == "node_flag":
-            return (-1, 0)
+            return (-2, 0)
+        elif name == "external_spec_selected":
+            return (0, 0)  # note out of order so this goes last
         else:
-            return (0, 0)
+            return (-1, 0)
 
     def build_specs(self, function_tuples):
         # Functions don't seem to be in particular order in output.  Sort
@@ -2405,12 +2422,6 @@ class SpecBuilder(object):
                 if isinstance(spec.version, spack.version.GitVersion):
                     spec.version.generate_git_lookup(spec.fullname)
 
-        # Add synthetic edges for externals that are extensions
-        for root in self._specs.values():
-            for dep in root.traverse():
-                if dep.external:
-                    dep.package.update_external_dependencies()
-
         return self._specs
 
 
diff --git a/lib/spack/spack/test/concretize.py b/lib/spack/spack/test/concretize.py
index b693b19787..3e816dc63e 100644
--- a/lib/spack/spack/test/concretize.py
+++ b/lib/spack/spack/test/concretize.py
@@ -14,6 +14,7 @@ import llnl.util.lang
 
 import spack.compilers
 import spack.concretize
+import spack.detection
 import spack.error
 import spack.hash_types as ht
 import spack.platforms
@@ -1940,7 +1941,9 @@ class TestConcretize(object):
             assert s.satisfies("target=%s" % spack.platforms.test.Test.front_end)
 
     def test_external_python_extensions_have_dependency(self):
-        """Test that python extensions have access to a python dependency"""
+        """Test that python extensions have access to a python dependency
+
+        when python is otherwise in the DAG"""
         external_conf = {
             "py-extension1": {
                 "buildable": False,
@@ -1953,3 +1956,106 @@ class TestConcretize(object):
 
         assert "python" in spec["py-extension1"]
         assert spec["python"] == spec["py-extension1"]["python"]
+
+    target = spack.platforms.test.Test.default
+
+    @pytest.mark.parametrize(
+        "python_spec",
+        [
+            "python@configured",
+            "python@configured platform=test",
+            "python@configured os=debian",
+            "python@configured target=%s" % target,
+        ],
+    )
+    def test_external_python_extension_find_dependency_from_config(self, python_spec):
+        fake_path = os.path.sep + "fake"
+
+        external_conf = {
+            "py-extension1": {
+                "buildable": False,
+                "externals": [{"spec": "py-extension1@2.0", "prefix": fake_path}],
+            },
+            "python": {
+                "externals": [{"spec": python_spec, "prefix": fake_path}],
+            },
+        }
+        spack.config.set("packages", external_conf)
+
+        spec = Spec("py-extension1").concretized()
+
+        assert "python" in spec["py-extension1"]
+        assert spec["python"].prefix == fake_path
+        # The spec is not equal to spack.spec.Spec("python@configured") because it gets a
+        # namespace and an external prefix before marking concrete
+        assert spec["python"].satisfies(python_spec)
+
+    def test_external_python_extension_find_dependency_from_installed(self, monkeypatch):
+        fake_path = os.path.sep + "fake"
+
+        external_conf = {
+            "py-extension1": {
+                "buildable": False,
+                "externals": [{"spec": "py-extension1@2.0", "prefix": fake_path}],
+            },
+            "python": {
+                "buildable": False,
+                "externals": [{"spec": "python@installed", "prefix": fake_path}],
+            },
+        }
+        spack.config.set("packages", external_conf)
+
+        # install python external
+        python = Spec("python").concretized()
+        monkeypatch.setattr(spack.store.db, "query", lambda x: [python])
+
+        # ensure that we can't be faking this by getting it from config
+        external_conf.pop("python")
+        spack.config.set("packages", external_conf)
+
+        spec = Spec("py-extension1").concretized()
+
+        assert "python" in spec["py-extension1"]
+        assert spec["python"].prefix == fake_path
+        # The spec is not equal to spack.spec.Spec("python@configured") because it gets a
+        # namespace and an external prefix before marking concrete
+        assert spec["python"].satisfies(python)
+
+    def test_external_python_extension_find_dependency_from_detection(self, monkeypatch):
+        """Test that python extensions have access to a python dependency
+
+        when python isn't otherwise in the DAG"""
+        python_spec = spack.spec.Spec("python@detected")
+        prefix = os.path.sep + "fake"
+
+        def find_fake_python(classes, path_hints):
+            return {"python": [spack.detection.DetectedPackage(python_spec, prefix=path_hints[0])]}
+
+        monkeypatch.setattr(spack.detection, "by_executable", find_fake_python)
+        external_conf = {
+            "py-extension1": {
+                "buildable": False,
+                "externals": [{"spec": "py-extension1@2.0", "prefix": "%s" % prefix}],
+            }
+        }
+        spack.config.set("packages", external_conf)
+
+        spec = Spec("py-extension1").concretized()
+
+        assert "python" in spec["py-extension1"]
+        assert spec["python"].prefix == prefix
+        assert spec["python"] == python_spec
+
+    def test_external_python_extension_find_unified_python(self):
+        """Test that python extensions use the same python as other specs in unified env"""
+        external_conf = {
+            "py-extension1": {
+                "buildable": False,
+                "externals": [{"spec": "py-extension1@2.0", "prefix": os.path.sep + "fake"}],
+            }
+        }
+        spack.config.set("packages", external_conf)
+
+        abstract_specs = [spack.spec.Spec(s) for s in ["py-extension1", "python"]]
+        specs = spack.concretize.concretize_specs_together(*abstract_specs)
+        assert specs[0]["python"] == specs[1]["python"]