Merge pull request #9775 from uranusjr/new-resolver-direct-url-with-extras

Correctly resolve requirement requested both as non-extra URL and non-URL with extras

6 files changed, 169 insertions, 71 deletions

diff --git a/news/8785.bugfix.rst b/news/8785.bugfix.rst
new file mode 100644
index 000000000..b84d8d8d5
--- /dev/null
+++ b/news/8785.bugfix.rst
@@ -0,0 +1,4 @@
+New resolver: When a requirement is requested both via a direct URL
+(``req @ URL``) and via version specifier with extras (``req[extra]``), the
+resolver will now be able to use the URL to correctly resolve the requirement
+with extras.
diff --git a/src/pip/_internal/resolution/resolvelib/candidates.py b/src/pip/_internal/resolution/resolvelib/candidates.py
index b3c3d019c..da516ad3c 100644
--- a/src/pip/_internal/resolution/resolvelib/candidates.py
+++ b/src/pip/_internal/resolution/resolvelib/candidates.py
@@ -33,6 +33,18 @@ BaseCandidate = Union[
 ]
 
 
+def as_base_candidate(candidate: Candidate) -> Optional[BaseCandidate]:
+    """The runtime version of BaseCandidate."""
+    base_candidate_classes = (
+        AlreadyInstalledCandidate,
+        EditableCandidate,
+        LinkCandidate,
+    )
+    if isinstance(candidate, base_candidate_classes):
+        return candidate
+    return None
+
+
 def make_install_req_from_link(link, template):
     # type: (Link, InstallRequirement) -> InstallRequirement
     assert not template.editable, "template is editable"
diff --git a/src/pip/_internal/resolution/resolvelib/factory.py b/src/pip/_internal/resolution/resolvelib/factory.py
index 2dcd83894..6e3f19518 100644
--- a/src/pip/_internal/resolution/resolvelib/factory.py
+++ b/src/pip/_internal/resolution/resolvelib/factory.py
@@ -1,3 +1,4 @@
+import contextlib
 import functools
 import logging
 from typing import (
@@ -16,6 +17,8 @@ from typing import (
     cast,
 )
 
+from pip._vendor.packaging.requirements import InvalidRequirement
+from pip._vendor.packaging.requirements import Requirement as PackagingRequirement
 from pip._vendor.packaging.specifiers import SpecifierSet
 from pip._vendor.packaging.utils import NormalizedName, canonicalize_name
 from pip._vendor.pkg_resources import Distribution
@@ -54,6 +57,7 @@ from .candidates import (
     ExtrasCandidate,
     LinkCandidate,
     RequiresPythonCandidate,
+    as_base_candidate,
 )
 from .found_candidates import FoundCandidates, IndexCandidateInfo
 from .requirements import (
@@ -123,6 +127,15 @@ class Factory:
         # type: () -> bool
         return self._force_reinstall
 
+    def _fail_if_link_is_unsupported_wheel(self, link: Link) -> None:
+        if not link.is_wheel:
+            return
+        wheel = Wheel(link.filename)
+        if wheel.supported(self._finder.target_python.get_tags()):
+            return
+        msg = f"{link.filename} is not a supported wheel on this platform."
+        raise UnsupportedWheel(msg)
+
     def _make_extras_candidate(self, base, extras):
         # type: (BaseCandidate, FrozenSet[str]) -> ExtrasCandidate
         cache_key = (id(base), extras)
@@ -275,6 +288,51 @@ class Factory:
             incompatible_ids,
         )
 
+    def _iter_explicit_candidates_from_base(
+        self,
+        base_requirements: Iterable[Requirement],
+        extras: FrozenSet[str],
+    ) -> Iterator[Candidate]:
+        """Produce explicit candidates from the base given an extra-ed package.
+
+        :param base_requirements: Requirements known to the resolver. The
+            requirements are guaranteed to not have extras.
+        :param extras: The extras to inject into the explicit requirements'
+            candidates.
+        """
+        for req in base_requirements:
+            lookup_cand, _ = req.get_candidate_lookup()
+            if lookup_cand is None:  # Not explicit.
+                continue
+            # We've stripped extras from the identifier, and should always
+            # get a BaseCandidate here, unless there's a bug elsewhere.
+            base_cand = as_base_candidate(lookup_cand)
+            assert base_cand is not None, "no extras here"
+            yield self._make_extras_candidate(base_cand, extras)
+
+    def _iter_candidates_from_constraints(
+        self,
+        identifier: str,
+        constraint: Constraint,
+        template: InstallRequirement,
+    ) -> Iterator[Candidate]:
+        """Produce explicit candidates from constraints.
+
+        This creates "fake" InstallRequirement objects that are basically clones
+        of what "should" be the template, but with original_link set to link.
+        """
+        for link in constraint.links:
+            self._fail_if_link_is_unsupported_wheel(link)
+            candidate = self._make_candidate_from_link(
+                link,
+                extras=frozenset(),
+                template=install_req_from_link_and_ireq(link, template),
+                name=canonicalize_name(identifier),
+                version=None,
+            )
+            if candidate:
+                yield candidate
+
     def find_candidates(
         self,
         identifier: str,
@@ -283,59 +341,48 @@ class Factory:
         constraint: Constraint,
         prefers_installed: bool,
     ) -> Iterable[Candidate]:
-
-        # Since we cache all the candidates, incompatibility identification
-        # can be made quicker by comparing only the id() values.
-        incompat_ids = {id(c) for c in incompatibilities.get(identifier, ())}
-
+        # Collect basic lookup information from the requirements.
         explicit_candidates = set()  # type: Set[Candidate]
         ireqs = []  # type: List[InstallRequirement]
         for req in requirements[identifier]:
             cand, ireq = req.get_candidate_lookup()
-            if cand is not None and id(cand) not in incompat_ids:
+            if cand is not None:
                 explicit_candidates.add(cand)
             if ireq is not None:
                 ireqs.append(ireq)
 
-        for link in constraint.links:
-            if not ireqs:
-                # If we hit this condition, then we cannot construct a candidate.
-                # However, if we hit this condition, then none of the requirements
-                # provided an ireq, so they must have provided an explicit candidate.
-                # In that case, either the candidate matches, in which case this loop
-                # doesn't need to do anything, or it doesn't, in which case there's
-                # nothing this loop can do to recover.
-                break
-            if link.is_wheel:
-                wheel = Wheel(link.filename)
-                # Check whether the provided wheel is compatible with the target
-                # platform.
-                if not wheel.supported(self._finder.target_python.get_tags()):
-                    # We are constrained to install a wheel that is incompatible with
-                    # the target architecture, so there are no valid candidates.
-                    # Return early, with no candidates.
-                    return ()
-            # Create a "fake" InstallRequirement that's basically a clone of
-            # what "should" be the template, but with original_link set to link.
-            # Using the given requirement is necessary for preserving hash
-            # requirements, but without the original_link, direct_url.json
-            # won't be created.
-            ireq = install_req_from_link_and_ireq(link, ireqs[0])
-            candidate = self._make_candidate_from_link(
-                link,
-                extras=frozenset(),
-                template=ireq,
-                name=canonicalize_name(ireq.name) if ireq.name else None,
-                version=None,
+        # If the current identifier contains extras, add explicit candidates
+        # from entries from extra-less identifier.
+        with contextlib.suppress(InvalidRequirement):
+            parsed_requirement = PackagingRequirement(identifier)
+            explicit_candidates.update(
+                self._iter_explicit_candidates_from_base(
+                    requirements.get(parsed_requirement.name, ()),
+                    frozenset(parsed_requirement.extras),
+                ),
             )
-            if candidate is None:
-                # _make_candidate_from_link returns None if the wheel fails to build.
-                # We are constrained to install this wheel, so there are no valid
-                # candidates.
-                # Return early, with no candidates.
+
+        # Add explicit candidates from constraints. We only do this if there are
+        # kown ireqs, which represent requirements not already explicit. If
+        # there are no ireqs, we're constraining already-explicit requirements,
+        # which is handled later when we return the explicit candidates.
+        if ireqs:
+            try:
+                explicit_candidates.update(
+                    self._iter_candidates_from_constraints(
+                        identifier,
+                        constraint,
+                        template=ireqs[0],
+                    ),
+                )
+            except UnsupportedWheel:
+                # If we're constrained to install a wheel incompatible with the
+                # target architecture, no candidates will ever be valid.
                 return ()
 
-            explicit_candidates.add(candidate)
+        # Since we cache all the candidates, incompatibility identification
+        # can be made quicker by comparing only the id() values.
+        incompat_ids = {id(c) for c in incompatibilities.get(identifier, ())}
 
         # If none of the requirements want an explicit candidate, we can ask
         # the finder for candidates.
@@ -351,7 +398,8 @@ class Factory:
         return (
             c
             for c in explicit_candidates
-            if constraint.is_satisfied_by(c)
+            if id(c) not in incompat_ids
+            and constraint.is_satisfied_by(c)
             and all(req.is_satisfied_by(c) for req in requirements[identifier])
         )
 
@@ -366,13 +414,7 @@ class Factory:
             return None
         if not ireq.link:
             return SpecifierRequirement(ireq)
-        if ireq.link.is_wheel:
-            wheel = Wheel(ireq.link.filename)
-            if not wheel.supported(self._finder.target_python.get_tags()):
-                msg = "{} is not a supported wheel on this platform.".format(
-                    wheel.filename,
-                )
-                raise UnsupportedWheel(msg)
+        self._fail_if_link_is_unsupported_wheel(ireq.link)
         cand = self._make_candidate_from_link(
             ireq.link,
             extras=frozenset(ireq.extras),
diff --git a/tests/functional/test_install_direct_url.py b/tests/functional/test_install_direct_url.py
index e28a7e9b5..baa5a3f2c 100644
--- a/tests/functional/test_install_direct_url.py
+++ b/tests/functional/test_install_direct_url.py
@@ -1,26 +1,12 @@
-import re
-
 import pytest
 
-from pip._internal.models.direct_url import DIRECT_URL_METADATA_NAME, DirectUrl
 from tests.lib import _create_test_package, path_to_url
-
-
-def _get_created_direct_url(result, pkg):
-    direct_url_metadata_re = re.compile(
-        pkg + r"-[\d\.]+\.dist-info." + DIRECT_URL_METADATA_NAME + r"$"
-    )
-    for filename in result.files_created:
-        if direct_url_metadata_re.search(filename):
-            direct_url_path = result.test_env.base_path / filename
-            with open(direct_url_path) as f:
-                return DirectUrl.from_json(f.read())
-    return None
+from tests.lib.direct_url import get_created_direct_url
 
 
 def test_install_find_links_no_direct_url(script, with_wheel):
     result = script.pip_install_local("simple")
-    assert not _get_created_direct_url(result, "simple")
+    assert not get_created_direct_url(result, "simple")
 
 
 def test_install_vcs_editable_no_direct_url(script, with_wheel):
@@ -29,7 +15,7 @@ def test_install_vcs_editable_no_direct_url(script, with_wheel):
     result = script.pip(*args)
     # legacy editable installs do not generate .dist-info,
     # hence no direct_url.json
-    assert not _get_created_direct_url(result, "testpkg")
+    assert not get_created_direct_url(result, "testpkg")
 
 
 def test_install_vcs_non_editable_direct_url(script, with_wheel):
@@ -37,7 +23,7 @@ def test_install_vcs_non_editable_direct_url(script, with_wheel):
     url = path_to_url(pkg_path)
     args = ["install", f"git+{url}#egg=testpkg"]
     result = script.pip(*args)
-    direct_url = _get_created_direct_url(result, "testpkg")
+    direct_url = get_created_direct_url(result, "testpkg")
     assert direct_url
     assert direct_url.url == url
     assert direct_url.info.vcs == "git"
@@ -47,7 +33,7 @@ def test_install_archive_direct_url(script, data, with_wheel):
     req = "simple @ " + path_to_url(data.packages / "simple-2.0.tar.gz")
     assert req.startswith("simple @ file://")
     result = script.pip("install", req)
-    assert _get_created_direct_url(result, "simple")
+    assert get_created_direct_url(result, "simple")
 
 
 @pytest.mark.network
@@ -59,7 +45,7 @@ def test_install_vcs_constraint_direct_url(script, with_wheel):
         "#egg=pip-test-package"
     )
     result = script.pip("install", "pip-test-package", "-c", constraints_file)
-    assert _get_created_direct_url(result, "pip_test_package")
+    assert get_created_direct_url(result, "pip_test_package")
 
 
 def test_install_vcs_constraint_direct_file_url(script, with_wheel):
@@ -68,4 +54,4 @@ def test_install_vcs_constraint_direct_file_url(script, with_wheel):
     constraints_file = script.scratch_path / "constraints.txt"
     constraints_file.write_text(f"git+{url}#egg=testpkg")
     result = script.pip("install", "testpkg", "-c", constraints_file)
-    assert _get_created_direct_url(result, "testpkg")
+    assert get_created_direct_url(result, "testpkg")
diff --git a/tests/functional/test_new_resolver.py b/tests/functional/test_new_resolver.py
index b2d3625a4..0938768a2 100644
--- a/tests/functional/test_new_resolver.py
+++ b/tests/functional/test_new_resolver.py
@@ -12,6 +12,7 @@ from tests.lib import (
     create_test_package_with_setup,
     path_to_url,
 )
+from tests.lib.direct_url import get_created_direct_url
 from tests.lib.path import Path
 from tests.lib.wheel import make_wheel
 
@@ -1788,3 +1789,41 @@ def test_new_resolver_avoids_incompatible_wheel_tags_in_constraint_url(
 
     assert_installed(script, base="0.1.0")
     assert_not_installed(script, "dep")
+
+
+def test_new_resolver_direct_url_with_extras(tmp_path, script):
+    pkg1 = create_basic_wheel_for_package(script, name="pkg1", version="1")
+    pkg2 = create_basic_wheel_for_package(
+        script,
+        name="pkg2",
+        version="1",
+        extras={"ext": ["pkg1"]},
+    )
+    pkg3 = create_basic_wheel_for_package(
+        script,
+        name="pkg3",
+        version="1",
+        depends=["pkg2[ext]"],
+    )
+
+    # Make pkg1 and pkg3 visible via --find-links, but not pkg2.
+    find_links = tmp_path.joinpath("find_links")
+    find_links.mkdir()
+    with open(pkg1, "rb") as f:
+        find_links.joinpath(pkg1.name).write_bytes(f.read())
+    with open(pkg3, "rb") as f:
+        find_links.joinpath(pkg3.name).write_bytes(f.read())
+
+    # Install with pkg2 only available with direct URL. The extra-ed direct
+    # URL pkg2 should be able to provide pkg2[ext] required by pkg3.
+    result = script.pip(
+        "install",
+        "--no-cache-dir", "--no-index",
+        "--find-links", str(find_links),
+        pkg2, "pkg3",
+    )
+
+    assert_installed(script, pkg1="1", pkg2="1", pkg3="1")
+    assert not get_created_direct_url(result, "pkg1")
+    assert get_created_direct_url(result, "pkg2")
+    assert not get_created_direct_url(result, "pkg3")
diff --git a/tests/lib/direct_url.py b/tests/lib/direct_url.py
new file mode 100644
index 000000000..497e10c6b
--- /dev/null
+++ b/tests/lib/direct_url.py
@@ -0,0 +1,15 @@
+import re
+
+from pip._internal.models.direct_url import DIRECT_URL_METADATA_NAME, DirectUrl
+
+
+def get_created_direct_url(result, pkg):
+    direct_url_metadata_re = re.compile(
+        pkg + r"-[\d\.]+\.dist-info." + DIRECT_URL_METADATA_NAME + r"$"
+    )
+    for filename in result.files_created:
+        if direct_url_metadata_re.search(filename):
+            direct_url_path = result.test_env.base_path / filename
+            with open(direct_url_path) as f:
+                return DirectUrl.from_json(f.read())
+    return None