Improve handling of constraints on requirements with extras

src/pip/_internal/req/constructors.py

@@ -8,14 +8,15 @@
 InstallRequirement.
 """
 
+import copy
 import logging
 import os
 import re
-from typing import Dict, List, Optional, Set, Tuple, Union
+from typing import Collection, Dict, List, Optional, Set, Tuple, Union
 
 from pip._vendor.packaging.markers import Marker
 from pip._vendor.packaging.requirements import InvalidRequirement, Requirement
-from pip._vendor.packaging.specifiers import Specifier
+from pip._vendor.packaging.specifiers import Specifier, SpecifierSet
 
 from pip._internal.exceptions import InstallationError
 from pip._internal.models.index import PyPI, TestPyPI
@@ -504,3 +505,45 @@ def install_req_from_link_and_ireq(
         config_settings=ireq.config_settings,
         user_supplied=ireq.user_supplied,
     )
+
+
+def install_req_drop_extras(ireq: InstallRequirement) -> InstallRequirement:
+    """
+    Creates a new InstallationRequirement using the given template but without
+    any extras. Sets the original requirement as the new one's parent
+    (comes_from).
+    """
+    req = Requirement(str(ireq.req))
+    req.extras = {}
+    return InstallRequirement(
+        req=req,
+        comes_from=ireq,
+        editable=ireq.editable,
+        link=ireq.link,
+        markers=ireq.markers,
+        use_pep517=ireq.use_pep517,
+        isolated=ireq.isolated,
+        global_options=ireq.global_options,
+        hash_options=ireq.hash_options,
+        constraint=ireq.constraint,
+        extras=[],
+        config_settings=ireq.config_settings,
+        user_supplied=ireq.user_supplied,
+        permit_editable_wheels=ireq.permit_editable_wheels,
+    )
+
+
+def install_req_extend_extras(
+    ireq: InstallRequirement,
+    extras: Collection[str],
+) -> InstallRequirement:
+    """
+    Returns a copy of an installation requirement with some additional extras.
+    Makes a shallow copy of the ireq object.
+    """
+    result = copy.copy(ireq)
+    req = Requirement(str(ireq.req))
+    req.extras.update(extras)
+    result.req = req
+    result.extras = {*ireq.extras, *extras}
+    return result

src/pip/_internal/resolution/resolvelib/candidates.py

@@ -240,7 +240,7 @@ def _prepare(self) -> BaseDistribution:
     def iter_dependencies(self, with_requires: bool) -> Iterable[Optional[Requirement]]:
         requires = self.dist.iter_dependencies() if with_requires else ()
         for r in requires:
-            yield self._factory.make_requirement_from_spec(str(r), self._ireq)
+            yield from self._factory.make_requirements_from_spec(str(r), self._ireq)
         yield self._factory.make_requires_python_requirement(self.dist.requires_python)
 
     def get_install_requirement(self) -> Optional[InstallRequirement]:
@@ -392,7 +392,7 @@ def iter_dependencies(self, with_requires: bool) -> Iterable[Optional[Requiremen
         if not with_requires:
             return
         for r in self.dist.iter_dependencies():
-            yield self._factory.make_requirement_from_spec(str(r), self._ireq)
+            yield from self._factory.make_requirements_from_spec(str(r), self._ireq)
 
     def get_install_requirement(self) -> Optional[InstallRequirement]:
         return None
@@ -427,9 +427,19 @@ def __init__(
         self,
         base: BaseCandidate,
         extras: FrozenSet[str],
+        ireq: Optional[InstallRequirement] = None,
     ) -> None:
+        """
+        :param ireq: the InstallRequirement that led to this candidate, if it
+            differs from the base's InstallRequirement. This will often be the
+            case in the sense that this candidate's requirement has the extras
+            while the base's does not. Unlike the InstallRequirement backed
+            candidates, this requirement is used solely for reporting purposes,
+            it does not do any leg work.
+        """
         self.base = base
         self.extras = extras
+        self._ireq = ireq
 
     def __str__(self) -> str:
         name, rest = str(self.base).split(" ", 1)
@@ -502,11 +512,11 @@ def iter_dependencies(self, with_requires: bool) -> Iterable[Optional[Requiremen
             )
 
         for r in self.base.dist.iter_dependencies(valid_extras):
-            requirement = factory.make_requirement_from_spec(
-                str(r), self.base._ireq, valid_extras
+            yield from factory.make_requirements_from_spec(
+                str(r),
+                self._ireq if self._ireq is not None else self.base._ireq,
+                valid_extras,
             )
-            if requirement:
-                yield requirement
 
     def get_install_requirement(self) -> Optional[InstallRequirement]:
         # We don't return anything here, because we always

src/pip/_internal/resolution/resolvelib/factory.py

@@ -138,13 +138,16 @@ def _fail_if_link_is_unsupported_wheel(self, link: Link) -> None:
         raise UnsupportedWheel(msg)
 
     def _make_extras_candidate(
-        self, base: BaseCandidate, extras: FrozenSet[str]
+        self,
+        base: BaseCandidate,
+        extras: FrozenSet[str],
+        ireq: Optional[InstallRequirement] = None,
     ) -> ExtrasCandidate:
         cache_key = (id(base), extras)
         try:
             candidate = self._extras_candidate_cache[cache_key]
         except KeyError:
-            candidate = ExtrasCandidate(base, extras)
+            candidate = ExtrasCandidate(base, extras, ireq=ireq)
             self._extras_candidate_cache[cache_key] = candidate
         return candidate
 
@@ -161,7 +164,7 @@ def _make_candidate_from_dist(
             self._installed_candidate_cache[dist.canonical_name] = base
         if not extras:
             return base
-        return self._make_extras_candidate(base, extras)
+        return self._make_extras_candidate(base, extras, ireq=template)
 
     def _make_candidate_from_link(
         self,
@@ -223,7 +226,7 @@ def _make_candidate_from_link(
 
         if not extras:
             return base
-        return self._make_extras_candidate(base, extras)
+        return self._make_extras_candidate(base, extras, ireq=template)
 
     def _iter_found_candidates(
         self,
@@ -385,16 +388,21 @@ def find_candidates(
             if ireq is not None:
                 ireqs.append(ireq)
 
-        # If the current identifier contains extras, add explicit candidates
-        # from entries from extra-less identifier.
+        # If the current identifier contains extras, add requires and explicit
+        # candidates from entries from extra-less identifier.
         with contextlib.suppress(InvalidRequirement):
             parsed_requirement = get_requirement(identifier)
-            explicit_candidates.update(
-                self._iter_explicit_candidates_from_base(
-                    requirements.get(parsed_requirement.name, ()),
-                    frozenset(parsed_requirement.extras),
-                ),
-            )
+            if parsed_requirement.name != identifier:
+                explicit_candidates.update(
+                    self._iter_explicit_candidates_from_base(
+                        requirements.get(parsed_requirement.name, ()),
+                        frozenset(parsed_requirement.extras),
+                    ),
+                )
+                for req in requirements.get(parsed_requirement.name, []):
+                    _, ireq = req.get_candidate_lookup()
+                    if ireq is not None:
+                        ireqs.append(ireq)
 
         # Add explicit candidates from constraints. We only do this if there are
         # known ireqs, which represent requirements not already explicit. If
@@ -437,18 +445,33 @@ def find_candidates(
             and all(req.is_satisfied_by(c) for req in requirements[identifier])
         )
 
-    def _make_requirement_from_install_req(
+    def _make_requirements_from_install_req(
         self, ireq: InstallRequirement, requested_extras: Iterable[str]
-    ) -> Optional[Requirement]:
+    ) -> list[Requirement]:
+        """
+        Returns requirement objects associated with the given InstallRequirement. In
+        most cases this will be a single object but the following special cases exist:
+            - the InstallRequirement has markers that do not apply -> result is empty
+            - the InstallRequirement has both a constraint and extras -> result is split
+                in two requirement objects: one with the constraint and one with the
+                extra. This allows centralized constraint handling for the base,
+                resulting in fewer candidate rejections.
+        """
         if not ireq.match_markers(requested_extras):
             logger.info(
                 "Ignoring %s: markers '%s' don't match your environment",
                 ireq.name,
                 ireq.markers,
             )
-            return None
+            return []
         if not ireq.link:
-            return SpecifierRequirement(ireq)
+            if ireq.extras and ireq.req.specifier:
+                return [
+                    SpecifierRequirement(ireq, drop_extras=True),
+                    SpecifierRequirement(ireq),
+                ]
+            else:
+                return [SpecifierRequirement(ireq)]
         self._fail_if_link_is_unsupported_wheel(ireq.link)
         cand = self._make_candidate_from_link(
             ireq.link,
@@ -466,8 +489,8 @@ def _make_requirement_from_install_req(
             # ResolutionImpossible eventually.
             if not ireq.name:
                 raise self._build_failures[ireq.link]
-            return UnsatisfiableRequirement(canonicalize_name(ireq.name))
-        return self.make_requirement_from_candidate(cand)
+            return [UnsatisfiableRequirement(canonicalize_name(ireq.name))]
+        return [self.make_requirement_from_candidate(cand)]
 
     def collect_root_requirements(
         self, root_ireqs: List[InstallRequirement]
@@ -488,30 +511,50 @@ def collect_root_requirements(
                 else:
                     collected.constraints[name] = Constraint.from_ireq(ireq)
             else:
-                req = self._make_requirement_from_install_req(
+                reqs = self._make_requirements_from_install_req(
                     ireq,
                     requested_extras=(),
                 )
-                if req is None:
+                if not reqs:
                     continue
-                if ireq.user_supplied and req.name not in collected.user_requested:
-                    collected.user_requested[req.name] = i
-                collected.requirements.append(req)
+
+                template = reqs[0]
+                if ireq.user_supplied and template.name not in collected.user_requested:
+                    collected.user_requested[template.name] = i
+                collected.requirements.extend(reqs)
+        # Put requirements with extras at the end of the root requires. This does not
+        # affect resolvelib's picking preference but it does affect its initial criteria
+        # population: by putting extras at the end we enable the candidate finder to
+        # present resolvelib with a smaller set of candidates to resolvelib, already
+        # taking into account any non-transient constraints on the associated base. This
+        # means resolvelib will have fewer candidates to visit and reject.
+        # Python's list sort is stable, meaning relative order is kept for objects with
+        # the same key.
+        collected.requirements.sort(key=lambda r: r.name != r.project_name)
         return collected
 
     def make_requirement_from_candidate(
         self, candidate: Candidate
     ) -> ExplicitRequirement:
         return ExplicitRequirement(candidate)
 
-    def make_requirement_from_spec(
+    def make_requirements_from_spec(
         self,
         specifier: str,
         comes_from: Optional[InstallRequirement],
         requested_extras: Iterable[str] = (),
-    ) -> Optional[Requirement]:
+    ) -> list[Requirement]:
+        """
+        Returns requirement objects associated with the given specifier. In most cases
+        this will be a single object but the following special cases exist:
+            - the specifier has markers that do not apply -> result is empty
+            - the specifier has both a constraint and extras -> result is split
+                in two requirement objects: one with the constraint and one with the
+                extra. This allows centralized constraint handling for the base,
+                resulting in fewer candidate rejections.
+        """
         ireq = self._make_install_req_from_spec(specifier, comes_from)
-        return self._make_requirement_from_install_req(ireq, requested_extras)
+        return self._make_requirements_from_install_req(ireq, requested_extras)
 
     def make_requires_python_requirement(
         self,

src/pip/_internal/resolution/resolvelib/requirements.py

@@ -2,6 +2,7 @@
 from pip._vendor.packaging.utils import NormalizedName, canonicalize_name
 
 from pip._internal.req.req_install import InstallRequirement
+from pip._internal.req.constructors import install_req_drop_extras
 
 from .base import Candidate, CandidateLookup, Requirement, format_name
 
@@ -40,13 +41,23 @@ def is_satisfied_by(self, candidate: Candidate) -> bool:
 
 
 class SpecifierRequirement(Requirement):
-    def __init__(self, ireq: InstallRequirement) -> None:
+    def __init__(
+        self,
+        ireq: InstallRequirement,
+        *,
+        drop_extras: bool = False,
+    ) -> None:
+        """
+        :param drop_extras: Ignore any extras that are part of the install requirement,
+            making this a requirement on the base only.
+        """
         assert ireq.link is None, "This is a link, not a specifier"
-        self._ireq = ireq
-        self._extras = frozenset(ireq.extras)
+        self._drop_extras: bool = drop_extras
+        self._ireq = ireq if not drop_extras else install_req_drop_extras(ireq)
+        self._extras = frozenset(self._ireq.extras)
 
     def __str__(self) -> str:
-        return str(self._ireq.req)
+        return str(self._ireq)
 
     def __repr__(self) -> str:
         return "{class_name}({requirement!r})".format(
@@ -61,7 +72,10 @@ def project_name(self) -> NormalizedName:
 
     @property
     def name(self) -> str:
-        return format_name(self.project_name, self._extras)
+        return format_name(
+            self.project_name,
+            self._extras,
+        )
 
     def format_for_error(self) -> str:
         # Convert comma-separated specifiers into "A, B, ..., F and G"

src/pip/_internal/resolution/resolvelib/resolver.py

@@ -1,3 +1,4 @@
+import contextlib
 import functools
 import logging
 import os
@@ -11,6 +12,7 @@
 from pip._internal.cache import WheelCache
 from pip._internal.index.package_finder import PackageFinder
 from pip._internal.operations.prepare import RequirementPreparer
+from pip._internal.req.constructors import install_req_extend_extras
 from pip._internal.req.req_install import InstallRequirement
 from pip._internal.req.req_set import RequirementSet
 from pip._internal.resolution.base import BaseResolver, InstallRequirementProvider
@@ -19,6 +21,7 @@
     PipDebuggingReporter,
     PipReporter,
 )
+from pip._internal.utils.packaging import get_requirement
 
 from .base import Candidate, Requirement
 from .factory import Factory
@@ -101,9 +104,19 @@ def resolve(
             raise error from e
 
         req_set = RequirementSet(check_supported_wheels=check_supported_wheels)
-        for candidate in result.mapping.values():
+        # sort to ensure base candidates come before candidates with extras
+        for candidate in sorted(result.mapping.values(), key=lambda c: c.name):
             ireq = candidate.get_install_requirement()
             if ireq is None:
+                if candidate.name != candidate.project_name:
+                    # extend existing req's extras
+                    with contextlib.suppress(KeyError):
+                        req = req_set.get_requirement(candidate.project_name)
+                        req_set.add_named_requirement(
+                            install_req_extend_extras(
+                                req, get_requirement(candidate.name).extras
+                            )
+                        )
                 continue
 
             # Check if there is already an installation under the same name,