Support custom setuptools & wheel versions.

When using non-vendored Pip (i.e.: specifying a custom `--pip-version`),
you can now override the versions of setuptools and wheel Pex bootstraps
for Pip with the existing `--extra-pip-requirement` option and the
override will be respected. Trying to override setuptools or wheel for
vendored Pip will raise an error however since its versions are
specialized to support all Pex resolve features under Python 2.7.

Fixes pex-tool#1895

pex/build_system/pep_517.py

@@ -45,7 +45,7 @@ def _default_build_system(
             resolved_reqs = set()  # type: Set[str]
             resolved_dists = []  # type: List[Distribution]
             if selected_pip_version is PipVersion.VENDORED:
-                requires = ["setuptools", selected_pip_version.wheel_requirement]
+                requires = ["setuptools", str(selected_pip_version.wheel_requirement)]
                 resolved_dists.extend(
                     Distribution.load(dist_location)
                     for dist_location in third_party.expose(
@@ -56,8 +56,8 @@ def _default_build_system(
                 extra_env.update(__PEX_UNVENDORED__="setuptools")
             else:
                 requires = [
-                    selected_pip_version.setuptools_requirement,
-                    selected_pip_version.wheel_requirement,
+                    str(selected_pip_version.setuptools_requirement),
+                    str(selected_pip_version.wheel_requirement),
                 ]
             unresolved = [
                 requirement for requirement in requires if requirement not in resolved_reqs

pex/pip/installation.py

@@ -5,6 +5,7 @@
 
 import hashlib
 import os
+from collections import OrderedDict
 from textwrap import dedent
 
 from pex import pex_warnings, third_party
@@ -13,6 +14,7 @@
 from pex.dist_metadata import Requirement
 from pex.interpreter import PythonInterpreter
 from pex.orderedset import OrderedSet
+from pex.pep_503 import ProjectName
 from pex.pex import PEX
 from pex.pex_bootstrapper import ensure_venv
 from pex.pip.tool import Pip, PipVenv
@@ -21,6 +23,7 @@
 from pex.result import Error, try_
 from pex.targets import LocalInterpreter, RequiresPythonError, Targets
 from pex.third_party import isolated
+from pex.tracer import TRACER
 from pex.typing import TYPE_CHECKING
 from pex.util import named_temporary_file
 from pex.variables import ENV
@@ -86,6 +89,11 @@ def _fingerprint(requirements):
     return hashlib.sha1("\n".join(sorted(map(str, requirements))).encode("utf-8")).hexdigest()
 
 
+_PIP_PROJECT_NAME = ProjectName("pip")
+_SETUPTOOLS_PROJECT_NAME = ProjectName("setuptools")
+_WHEEL_PROJECT_NAME = ProjectName("wheel")
+
+
 def _vendored_installation(
     interpreter=None,  # type: Optional[PythonInterpreter]
     resolver=None,  # type: Optional[Resolver]
@@ -114,6 +122,33 @@ def expose_vendored():
             )
         )
 
+    # Ensure user-specified extra requirements do not override vendored Pip or its setuptools and
+    # wheel dependencies. These are arranged just so with some patching to Pip and setuptools as
+    # well as a low enough standard wheel version to support Python 2.7.
+    for extra_req in extra_requirements:
+        if _PIP_PROJECT_NAME == extra_req.project_name:
+            raise ValueError(
+                "An `--extra-pip-requirement` cannot be used to override the Pip version; use "
+                "`--pip-version` to select a supported Pip version instead. "
+                "Given: {pip_req}".format(pip_req=extra_req)
+            )
+        if _SETUPTOOLS_PROJECT_NAME == extra_req.project_name:
+            raise ValueError(
+                "An `--extra-pip-requirement` cannot be used to override the setuptools version "
+                "for vendored Pip. If you need a custom setuptools you need to use `--pip-version` "
+                "to select a non-vendored Pip version. Given: {setuptools_req}".format(
+                    setuptools_req=extra_req
+                )
+            )
+        if _WHEEL_PROJECT_NAME == extra_req.project_name:
+            raise ValueError(
+                "An `--extra-pip-requirement` cannot be used to override the wheel version for "
+                "vendored Pip. If you need a custom wheel version you need to use `--pip-version` "
+                "to select a non-vendored Pip version. Given: {wheel_req}".format(
+                    wheel_req=extra_req
+                )
+            )
+
     # 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
@@ -156,9 +191,9 @@ def bootstrap_pip():
         )
 
         for req in version.requirements:
-            project_name = Requirement.parse(req).name
+            project_name = req.name
             target_dir = os.path.join(chroot, "reqs", project_name)
-            venv.interpreter.execute(["-m", "pip", "install", "--target", target_dir, req])
+            venv.interpreter.execute(["-m", "pip", "install", "--target", target_dir, str(req)])
             yield target_dir
 
     return bootstrap_pip
@@ -189,20 +224,40 @@ def _resolved_installation(
             fingerprint=_fingerprint(extra_requirements),
         )
 
-    requirements = list(version.requirements)
-    requirements.extend(map(str, extra_requirements))
+    requirements_by_project_name = OrderedDict(
+        (req.project_name, str(req)) for req in version.requirements
+    )
+
+    # Allow user-specified extra requirements to override Pip requirements (setuptools and wheel).
+    for extra_req in extra_requirements:
+        if _PIP_PROJECT_NAME == extra_req.project_name:
+            raise ValueError(
+                "An `--extra-pip-requirement` cannot be used to override the Pip version; use "
+                "`--pip-version` to select a supported Pip version instead. "
+                "Given: {pip_req}".format(pip_req=extra_req)
+            )
+        existing_req = requirements_by_project_name.get(extra_req.project_name)
+        if existing_req:
+            TRACER.log(
+                "Overriding `--pip-version {pip_version}` requirement of {existing_req} with "
+                "user-specified requirement {extra_req}".format(
+                    pip_version=version.version, existing_req=existing_req, extra_req=extra_req
+                )
+            )
+        requirements_by_project_name[extra_req.project_name] = str(extra_req)
+
     if not resolver:
         raise ValueError(
             "A resolver is required to install {requirements} for Pip {version}: {reqs}".format(
-                requirements=pluralize(requirements, "requirement"),
+                requirements=pluralize(requirements_by_project_name, "requirement"),
                 version=version,
-                reqs=" ".join(map(str, extra_requirements)),
+                reqs=" ".join(requirements_by_project_name.values()),
             )
         )
 
     def resolve_distribution_locations():
         for resolved_distribution in resolver.resolve_requirements(
-            requirements=requirements,
+            requirements=requirements_by_project_name.values(),
             targets=targets,
             pip_version=bootstrap_pip_version,
             extra_resolver_requirements=(),

pex/pip/tool.py

@@ -293,6 +293,11 @@ class Pip(object):
     version = attr.ib()  # type: PipVersionValue
     _pip_cache = attr.ib()  # type: str
 
+    @property
+    def venv_dir(self):
+        # type: () -> str
+        return self._pip.venv_dir
+
     @staticmethod
     def _calculate_resolver_version(package_index_configuration=None):
         # type: (Optional[PackageIndexConfiguration]) -> ResolverVersion.Value
@@ -643,7 +648,7 @@ def _ensure_wheel_installed(self, package_index_configuration=None):
             if not atomic_dir.is_finalized():
                 self.spawn_download_distributions(
                     download_dir=atomic_dir.work_dir,
-                    requirements=[self.version.wheel_requirement],
+                    requirements=[str(self.version.wheel_requirement)],
                     package_index_configuration=package_index_configuration,
                     build_configuration=BuildConfiguration.create(allow_builds=False),
                 ).wait()

pex/pip/version.py

@@ -58,8 +58,8 @@ def to_requirement(
             project_name,  # type: str
             project_version=None,  # type: Optional[str]
         ):
-            # type: (...) -> str
-            return (
+            # type: (...) -> Requirement
+            return Requirement.parse(
                 "{project_name}=={project_version}".format(
                     project_name=project_name, project_version=project_version
                 )
@@ -68,19 +68,21 @@ def to_requirement(
             )
 
         self.version = Version(version)
-        self.requirement = requirement or to_requirement("pip", version)
+        self.requirement = (
+            Requirement.parse(requirement) if requirement else to_requirement("pip", version)
+        )
         self.setuptools_requirement = to_requirement("setuptools", setuptools_version)
         self.wheel_requirement = to_requirement("wheel", wheel_version)
         self.requires_python = SpecifierSet(requires_python) if requires_python else None
         self.hidden = hidden
 
     @property
     def requirements(self):
-        # type: () -> Iterable[str]
+        # type: () -> Iterable[Requirement]
         return self.requirement, self.setuptools_requirement, self.wheel_requirement
 
     def requires_python_applies(self, target=None):
-        # type: (Optional[Union[Version,Target]]) -> bool
+        # type: (Optional[Union[Version, Target]]) -> bool
         if not self.requires_python:
             return True
 
@@ -89,10 +91,7 @@ def requires_python_applies(self, target=None):
 
         return LocalInterpreter.create(
             interpreter=target.get_interpreter() if target else None
-        ).requires_python_applies(
-            requires_python=self.requires_python,
-            source=Requirement.parse(self.requirement),
-        )
+        ).requires_python_applies(requires_python=self.requires_python, source=self.requirement)
 
     def __lt__(self, other):
         if not isinstance(other, PipVersionValue):
@@ -180,10 +179,6 @@ def values(cls):
         requires_python="<3.12",
     )
 
-    # TODO(John Sirois): Expose setuptools and wheel version flags - these don't affect
-    #  Pex; so we should allow folks to experiment with upgrade easily:
-    #  https://github.com/pex-tool/pex/issues/1895
-
     v22_2_2 = PipVersionValue(
         version="22.2.2",
         setuptools_version="65.3.0",

tests/integration/cli/commands/test_lock_sync.py

@@ -89,8 +89,8 @@ def host_requirements(*requirements):
             # itself if needed.
             host_requirements(
                 "cowsay==5.0.0",
-                pip_version.setuptools_requirement,
-                pip_version.wheel_requirement,
+                str(pip_version.setuptools_requirement),
+                str(pip_version.wheel_requirement),
             )
             find_links_repo.make_sdist("spam", version="1")
             find_links_repo.make_wheel("spam", version="1")

tests/integration/cli/commands/test_lock_update.py

@@ -80,8 +80,8 @@ def find_links(
     repository_pex = os.path.join(str(tmpdir), "repository.pex")
     run_pex_command(
         args=[
-            pip_version.setuptools_requirement,
-            pip_version.wheel_requirement,
+            str(pip_version.setuptools_requirement),
+            str(pip_version.wheel_requirement),
             "--include-tools",
             "-o",
             repository_pex,

tests/integration/test_issue_2343.py

@@ -41,8 +41,8 @@ def find_links(shared_integration_test_tmpdir):
             result = find_links_repo.resolver.resolve_requirements(
                 [
                     "ansicolors==1.1.8",
-                    pip_version.setuptools_requirement,
-                    pip_version.wheel_requirement,
+                    str(pip_version.setuptools_requirement),
+                    str(pip_version.wheel_requirement),
                 ],
                 result_type=InstallableType.WHEEL_FILE,
             )

tests/integration/test_keyring_support.py

@@ -201,7 +201,7 @@ def download_pip_requirements(
     extra_requirements=(),  # type: Iterable[str]
 ):
     # type: (...) -> None
-    requirements = list(pip_version.requirements)
+    requirements = list(map(str, pip_version.requirements))
     requirements.extend(extra_requirements)
     get_pip(resolver=ConfiguredResolver.version(pip_version)).spawn_download_distributions(
         download_dir=download_dir, requirements=requirements

tests/test_bdist_pex.py

@@ -219,7 +219,7 @@ def test_unwriteable_contents():
         wheels.extend(
             fingerprinted_dist.distribution.location
             for fingerprinted_dist in resolve(
-                requirements=[PipVersion.VENDORED.wheel_requirement],
+                requirements=[str(PipVersion.VENDORED.wheel_requirement)],
                 result_type=InstallableType.WHEEL_FILE,
             ).distributions
         )