Add support for --extra-pip-requirement. (#2461)

You can now inject extra requirements into the isolated Pip PEXes Pex
uses to resolve distributions. The motivating use case for this is to
use the feature Pip 23.1 introduced for forcing `--keyring-provider
import`.

Pex already supported using a combination of the following to force
non-interactive use of the keyring:
1. A `keyring` script installation that was on the `PATH`
2. A `--pip-version` 23.1 or newer.
3. Specifying `--use-pip-config` to pass `--keyring-provider subprocess`
   to Pip.

You could not force `--keyring-provider import` though since the Pips
Pex uses are themselves hermetic PEXes without access to extra installed
keyring requirements elsewhere on the system. With
`--extra-pip-requirement` you can do this with the primary benefit over
`--keyring-provider subprocess` being that you do not need to add the
username to index URLs. This is ultimately because the keyring CLI
requires username whereas the API does not; but see
https://pip.pypa.io/en/stable/topics/authentication/#keyring-support for
more information.

CHANGES.md

@@ -1,5 +1,31 @@
 # Release Notes
 
+## 2.10.0
+
+This release adds support for injecting requirements into the isolated
+Pip PEXes Pex uses to resolve distributions. The motivating use case
+for this is to use the feature Pip 23.1 introduced for forcing
+`--keyring-provider import`.
+
+Pex already supported using a combination of the following to force
+non-interactive use of the keyring:
+1. A `keyring` script installation that was on the `PATH`
+2. A `--pip-version` 23.1 or newer.
+3. Specifying `--use-pip-config` to pass `--keyring-provider subprocess`
+   to Pip.
+
+You could not force `--keyring-provider import` though, since the Pips
+Pex uses are themselves hermetic PEXes without access to extra
+installed  keyring requirements elsewhere on the system. With
+`--extra-pip-requirement` you can now do this with the primary benefit
+over `--keyring-provider subprocess` being that you do not need to add
+the username to index URLs. This is ultimately because the keyring CLI
+requires username whereas the API does not; but see
+https://pip.pypa.io/en/stable/topics/authentication/#keyring-support for
+more information.
+
+* Add support for `--extra-pip-requirement`. (#2461)
+
 ## 2.9.0
 
 This release adds support for Pip 24.1.2.

pex/cli/commands/lock.py

@@ -1668,6 +1668,7 @@ def _sync(self):
                 max_parallel_jobs=pip_configuration.max_jobs,
                 pip_version=pip_configuration.version,
                 use_pip_config=pip_configuration.use_pip_config,
+                extra_pip_requirements=pip_configuration.extra_requirements,
                 result_type=InstallableType.INSTALLED_WHEEL_CHROOT,
             )
         )

pex/pip/installation.py

@@ -3,12 +3,13 @@
 
 from __future__ import absolute_import
 
+import hashlib
 import os
 from textwrap import dedent
 
 from pex import pex_warnings, third_party
 from pex.atomic_directory import atomic_directory
-from pex.common import safe_mkdtemp
+from pex.common import pluralize, safe_mkdtemp
 from pex.dist_metadata import Requirement
 from pex.interpreter import PythonInterpreter
 from pex.orderedset import OrderedSet
@@ -26,7 +27,7 @@
 from pex.venv.virtualenv import InstallationChoice, Virtualenv
 
 if TYPE_CHECKING:
-    from typing import Callable, Dict, Iterator, Optional, Union
+    from typing import Callable, Dict, Iterator, Optional, Tuple, Union
 
     import attr  # vendor:skip
 else:
@@ -36,13 +37,14 @@
 def _pip_installation(
     version,  # type: PipVersionValue
     iter_distribution_locations,  # type: Callable[[], Iterator[str]]
+    fingerprint,  # type: str
     interpreter=None,  # type: Optional[PythonInterpreter]
 ):
     # type: (...) -> Pip
     pip_root = os.path.join(ENV.PEX_ROOT, "pip", str(version))
     path = os.path.join(pip_root, "pip.pex")
     pip_interpreter = interpreter or PythonInterpreter.get()
-    pip_pex_path = os.path.join(path, isolated().pex_hash)
+    pip_pex_path = os.path.join(path, isolated().pex_hash, fingerprint)
     with atomic_directory(pip_pex_path) as chroot:
         if not chroot.is_finalized():
             from pex.pex_builder import PEXBuilder
@@ -77,15 +79,63 @@ def _pip_installation(
     return Pip(pip=pip_venv, version=version, pip_cache=pip_cache)
 
 
-def _vendored_installation(interpreter=None):
-    # type: (Optional[PythonInterpreter]) -> Pip
+def _fingerprint(requirements):
+    # type: (Tuple[Requirement, ...]) -> str
+    if not requirements:
+        return "no-extra-requirements"
+    return hashlib.sha1("\n".join(sorted(map(str, requirements))).encode("utf-8")).hexdigest()
+
+
+def _vendored_installation(
+    interpreter=None,  # type: Optional[PythonInterpreter]
+    resolver=None,  # type: Optional[Resolver]
+    extra_requirements=(),  # type: Tuple[Requirement, ...]
+):
+    # type: (...) -> Pip
+
+    def expose_vendored():
+        # type: () -> Iterator[str]
+        return third_party.expose(("pip", "setuptools"), interpreter=interpreter)
+
+    if not extra_requirements:
+        return _pip_installation(
+            version=PipVersion.VENDORED,
+            iter_distribution_locations=expose_vendored,
+            interpreter=interpreter,
+            fingerprint=_fingerprint(extra_requirements),
+        )
+
+    if not resolver:
+        raise ValueError(
+            "A resolver is required to install extra {requirements} for vendored Pip: "
+            "{extra_requirements}".format(
+                requirements=pluralize(extra_requirements, "requirement"),
+                extra_requirements=" ".join(map(str, extra_requirements)),
+            )
+        )
+
+    # This indirection works around MyPy type inference failing to see that
+    # `iter_distribution_locations` is only successfully defined when resolve is not None.
+    extra_requirement_resolver = resolver
+
+    def iter_distribution_locations():
+        # type: () -> Iterator[str]
+        for location in expose_vendored():
+            yield location
+
+        for resolved_distribution in extra_requirement_resolver.resolve_requirements(
+            requirements=tuple(map(str, extra_requirements)),
+            targets=Targets.from_target(LocalInterpreter.create(interpreter)),
+            pip_version=PipVersion.VENDORED,
+            extra_resolver_requirements=(),
+        ).distributions:
+            yield resolved_distribution.distribution.location
 
     return _pip_installation(
         version=PipVersion.VENDORED,
-        iter_distribution_locations=lambda: third_party.expose(
-            ("pip", "setuptools"), interpreter=interpreter
-        ),
+        iter_distribution_locations=iter_distribution_locations,
         interpreter=interpreter,
+        fingerprint=_fingerprint(extra_requirements),
     )
 
 
@@ -102,7 +152,7 @@ def bootstrap_pip():
         venv = Virtualenv.create(
             venv_dir=os.path.join(chroot, "pip"),
             interpreter=interpreter,
-            install_pip=InstallationChoice.UPGRADED,
+            install_pip=InstallationChoice.YES,
         )
 
         for req in version.requirements:
@@ -118,6 +168,7 @@ def _resolved_installation(
     version,  # type: PipVersionValue
     resolver=None,  # type: Optional[Resolver]
     interpreter=None,  # type: Optional[PythonInterpreter]
+    extra_requirements=(),  # type: Tuple[Requirement, ...]
 ):
     # type: (...) -> Pip
     targets = Targets.from_target(LocalInterpreter.create(interpreter))
@@ -130,39 +181,47 @@ def _resolved_installation(
             warn=False,
         )
     )
-    if bootstrap_pip_version is not PipVersion.VENDORED:
+    if bootstrap_pip_version is not PipVersion.VENDORED and not extra_requirements:
         return _pip_installation(
             version=version,
             iter_distribution_locations=_bootstrap_pip(version, interpreter=interpreter),
             interpreter=interpreter,
+            fingerprint=_fingerprint(extra_requirements),
         )
 
-    if resolver is None:
+    requirements = list(version.requirements)
+    requirements.extend(map(str, extra_requirements))
+    if not resolver:
         raise ValueError(
-            "A resolver is required to install {requirement}".format(
-                requirement=version.requirement
+            "A resolver is required to install {requirements} for Pip {version}: {reqs}".format(
+                requirements=pluralize(requirements, "requirement"),
+                version=version,
+                reqs=" ".join(map(str, extra_requirements)),
             )
         )
 
     def resolve_distribution_locations():
         for resolved_distribution in resolver.resolve_requirements(
-            requirements=version.requirements,
+            requirements=requirements,
             targets=targets,
-            pip_version=PipVersion.VENDORED,
+            pip_version=bootstrap_pip_version,
+            extra_resolver_requirements=(),
         ).distributions:
             yield resolved_distribution.distribution.location
 
     return _pip_installation(
         version=version,
         iter_distribution_locations=resolve_distribution_locations,
         interpreter=interpreter,
+        fingerprint=_fingerprint(extra_requirements),
     )
 
 
 @attr.s(frozen=True)
 class PipInstallation(object):
     interpreter = attr.ib()  # type: PythonInterpreter
     version = attr.ib()  # type: PipVersionValue
+    extra_requirements = attr.ib()  # type: Tuple[Requirement, ...]
 
     def check_python_applies(self):
         # type: () -> None
@@ -257,6 +316,7 @@ def get_pip(
     interpreter=None,
     version=None,  # type: Optional[PipVersionValue]
     resolver=None,  # type: Optional[Resolver]
+    extra_requirements=(),  # type: Tuple[Requirement, ...]
 ):
     # type: (...) -> Pip
     """Returns a lazily instantiated global Pip object that is safe for un-coordinated use."""
@@ -280,15 +340,23 @@ def get_pip(
     installation = PipInstallation(
         interpreter=interpreter or PythonInterpreter.get(),
         version=calculated_version,
+        extra_requirements=extra_requirements,
     )
     pip = _PIP.get(installation)
     if pip is None:
         installation.check_python_applies()
         if installation.version is PipVersion.VENDORED:
-            pip = _vendored_installation(interpreter=interpreter)
+            pip = _vendored_installation(
+                interpreter=interpreter,
+                resolver=resolver,
+                extra_requirements=installation.extra_requirements,
+            )
         else:
             pip = _resolved_installation(
-                version=installation.version, resolver=resolver, interpreter=interpreter
+                version=installation.version,
+                resolver=resolver,
+                interpreter=interpreter,
+                extra_requirements=installation.extra_requirements,
             )
         _PIP[installation] = pip
     return pip

pex/pip/tool.py

@@ -19,6 +19,7 @@
 from pex.common import safe_mkdir, safe_mkdtemp
 from pex.compatibility import get_stderr_bytes_buffer, shlex_quote, urlparse
 from pex.dependency_configuration import DependencyConfiguration
+from pex.dist_metadata import Requirement
 from pex.interpreter import PythonInterpreter
 from pex.jobs import Job
 from pex.network_configuration import NetworkConfiguration
@@ -148,6 +149,7 @@ def create(
         network_configuration=None,  # type: Optional[NetworkConfiguration]
         password_entries=(),  # type: Iterable[PasswordEntry]
         use_pip_config=False,  # type: bool
+        extra_pip_requirements=(),  # type: Tuple[Requirement, ...]
     ):
         # type: (...) -> PackageIndexConfiguration
         resolver_version = resolver_version or ResolverVersion.default(pip_version)
@@ -169,6 +171,7 @@ def create(
                 network_configuration=network_configuration, use_pip_config=use_pip_config
             ),
             use_pip_config=use_pip_config,
+            extra_pip_requirements=extra_pip_requirements,
             password_entries=password_entries,
         )
 
@@ -181,6 +184,7 @@ def __init__(
         use_pip_config,  # type: bool
         password_entries=(),  # type: Iterable[PasswordEntry]
         pip_version=None,  # type: Optional[PipVersionValue]
+        extra_pip_requirements=(),  # type: Tuple[Requirement, ...]
     ):
         # type: (...) -> None
         self.resolver_version = resolver_version  # type: ResolverVersion.Value
@@ -190,6 +194,7 @@ def __init__(
         self.use_pip_config = use_pip_config  # type: bool
         self.password_entries = password_entries  # type: Iterable[PasswordEntry]
         self.pip_version = pip_version  # type: Optional[PipVersionValue]
+        self.extra_pip_requirements = extra_pip_requirements  # type: Tuple[Requirement, ...]
 
 
 if TYPE_CHECKING:

pex/pip/version.py

@@ -3,6 +3,7 @@
 
 from __future__ import absolute_import
 
+import functools
 import os
 import sys
 from textwrap import dedent
@@ -18,6 +19,7 @@
     from typing import Iterable, Optional, Tuple, Union
 
 
+@functools.total_ordering
 class PipVersionValue(Enum.Value):
     @classmethod
     def overridden(cls):
@@ -87,6 +89,11 @@ def requires_python_applies(self, target=None):
             source=Requirement.parse(self.requirement),
         )
 
+    def __lt__(self, other):
+        if not isinstance(other, PipVersionValue):
+            return NotImplemented
+        return self.version < other.version
+
 
 class LatestPipVersion(object):
     def __get__(self, obj, objtype=None):

pex/resolve/configured_resolve.py

@@ -59,6 +59,7 @@ def resolve(
                     max_parallel_jobs=pip_configuration.max_jobs,
                     pip_version=lock.pip_version,
                     use_pip_config=pip_configuration.use_pip_config,
+                    extra_pip_requirements=pip_configuration.extra_requirements,
                     result_type=result_type,
                     dependency_configuration=dependency_configuration,
                 )
@@ -103,6 +104,7 @@ def resolve(
                 pip_version=resolver_configuration.version,
                 resolver=ConfiguredResolver(pip_configuration=resolver_configuration),
                 use_pip_config=resolver_configuration.use_pip_config,
+                extra_pip_requirements=resolver_configuration.extra_requirements,
                 result_type=result_type,
                 dependency_configuration=dependency_configuration,
             )

pex/resolve/configured_resolver.py

@@ -4,6 +4,7 @@
 from __future__ import absolute_import
 
 from pex import resolver
+from pex.dist_metadata import Requirement
 from pex.pep_427 import InstallableType
 from pex.pip.version import PipVersion, PipVersionValue
 from pex.resolve import lock_resolver
@@ -15,7 +16,7 @@
 from pex.typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Iterable, Optional
+    from typing import Iterable, Optional, Tuple
 
     import attr  # vendor:skip
 else:
@@ -70,6 +71,7 @@ def resolve_lock(
                 max_parallel_jobs=self.pip_configuration.max_jobs,
                 pip_version=pip_version or self.pip_configuration.version,
                 use_pip_config=self.pip_configuration.use_pip_config,
+                extra_pip_requirements=self.pip_configuration.extra_requirements,
                 result_type=result_type,
             )
         )
@@ -80,6 +82,7 @@ def resolve_requirements(
         targets=Targets(),  # type: Targets
         pip_version=None,  # type: Optional[PipVersionValue]
         transitive=None,  # type: Optional[bool]
+        extra_resolver_requirements=None,  # type: Optional[Tuple[Requirement, ...]]
         result_type=InstallableType.INSTALLED_WHEEL_CHROOT,  # type: InstallableType.Value
     ):
         # type: (...) -> ResolveResult
@@ -100,5 +103,8 @@ def resolve_requirements(
             pip_version=pip_version or self.pip_configuration.version,
             resolver=self,
             use_pip_config=self.pip_configuration.use_pip_config,
+            extra_pip_requirements=extra_resolver_requirements
+            if extra_resolver_requirements is not None
+            else self.pip_configuration.extra_requirements,
             result_type=result_type,
         )

pex/resolve/downloads.py

@@ -65,6 +65,7 @@ def _pip(self):
             interpreter=self.target.get_interpreter(),
             version=self.package_index_configuration.pip_version,
             resolver=self.resolver,
+            extra_requirements=self.package_index_configuration.extra_pip_requirements,
         )
 
     @staticmethod