Repository: pex-tool/pex

Issue:
feature(pex.tools): prefer a pex-file-provided package versions when `--pip` flag is also provided
Given a `reqs.pex` file that includes a version of `pip`, `pex-tools reqs.pex --pip` should prefer the versions of packages included in the pex.

This may mean that pex has to downgrade the version of pip (and possibly setuptools, which pip depends on) to match whatever is included in the pex file.

Today, pex creates the virtualenv using the appropriate commands to include `pip`. Then it installs the requirements from the pex file, but any packages that already exist in the virtualenv, like pip, do not get downgraded to the version in the pex file. This feature request should reverse that.

pex-tools should report a warning about the collision between the pex file and the `--pip` flag (since the `--pip` flag is the only reason the virtualenv will have contents before installing things from the pex file). That warning should be silenced by `--collisions-ok`.

### Background: `pants`

When running `./pants export ...`, pants uses pex-tools `venv --pip --collisions-ok` for mutable virtualenvs.
https://github.com/pantsbuild/pants/blob/8fef19f9050d42c949edf781125cdebd6851fc57/src/python/pants/backend/python/goals/export.py#L239-L253
(that link has the following, which I'm including here since github doesn't include cross-repo snippets):
```python
                    complete_pex_env.create_argv(
                        os.path.join(tmpdir_under_digest_root, pex_pex.exe),
                        *[
                            os.path.join(tmpdir_under_digest_root, requirements_pex.name),
                            "venv",
                            "--pip",
                            "--collisions-ok",
                            output_path,
                        ],
                    ),
                    {
                        **complete_pex_env.environment_dict(python=requirements_pex.python),
                        "PEX_MODULE": "pex.tools",
                    },
                ),
```

For StackStorm, we depend on pip+setuptools at runtime, and I want to make sure that the locked version is always used in the venv.

### Slack

This request follows on a discussion in slack:
- [slack question](https://pantsbuild.slack.com/archives/C046T6T9U/p1679940252822339) ([linen mirror](https://chat.pantsbuild.org/t/10116580/u04s45aha-with-pex-tools-1-reqs-pex-venv-pip-some-path-is-th#04e53c65-b5a2-45e9-bef4-77e052c8e920))
- [slack conversation](https://pantsbuild.slack.com/archives/C046T6T9U/p1679940303328629) ([linen mirror](https://chat.pantsbuild.org/t/10116582/there-is-not-you-d-need-to-run-pip-after-creating-the-venv-t#aa48bee8-5a2b-4c32-92e9-599544c57d5e))


Discussion:

Resolving patch:
diff --git a/pex/tools/commands/venv.py b/pex/tools/commands/venv.py
--- a/pex/tools/commands/venv.py
+++ b/pex/tools/commands/venv.py
@@ -6,15 +6,22 @@
 import errno
 import logging
 import os
+import subprocess
 from argparse import ArgumentParser
+from collections import OrderedDict
+from subprocess import CalledProcessError
 
 from pex import pex_warnings
 from pex.common import safe_delete, safe_rmtree
+from pex.dist_metadata import Distribution
 from pex.enum import Enum
 from pex.executor import Executor
+from pex.pep_440 import Version
+from pex.pep_503 import ProjectName
 from pex.pex import PEX
-from pex.result import Error, Ok, Result
+from pex.result import Error, Ok, Result, try_
 from pex.tools.command import PEXCommand
+from pex.tracer import TRACER
 from pex.typing import TYPE_CHECKING
 from pex.venv.bin_path import BinPath
 from pex.venv.install_scope import InstallScope
@@ -22,7 +29,7 @@
 from pex.venv.virtualenv import PipUnavailableError, Virtualenv
 
 if TYPE_CHECKING:
-    from typing import Optional
+    from typing import Iterable, Optional, Union
 
     import attr  # vendor:skip
 else:
@@ -73,6 +80,104 @@ def save(self, install_scope):
             fp.write(str(install_scope))
 
 
+def find_dist(
+    project_name,  # type: ProjectName
+    dists,  # type: Iterable[Distribution]
+):
+    # type: (...) -> Optional[Version]
+    for dist in dists:
+        if project_name == dist.metadata.project_name:
+            return dist.metadata.version
+    return None
+
+
+_PIP = ProjectName("pip")
+_SETUPTOOLS = ProjectName("setuptools")
+
+
+def ensure_pip_installed(
+    venv,  # type: Virtualenv
+    pex,  # type: PEX
+    scope,  # type: InstallScope.Value
+    collisions_ok,  # type: bool
+):
+    # type: (...) -> Union[Version, Error]
+
+    venv_pip_version = find_dist(_PIP, venv.iter_distributions())
+    if venv_pip_version:
+        TRACER.log(
+            "The venv at {venv_dir} already has Pip {version} installed.".format(
+                venv_dir=venv.venv_dir, version=venv_pip_version
+            )
+        )
+    else:
+        try:
+            venv.install_pip()
+        except PipUnavailableError as e:
+            return Error(
+                "The virtual environment was successfully created, but Pip was not "
+                "installed:\n{}".format(e)
+            )
+        venv_pip_version = find_dist(_PIP, venv.iter_distributions())
+        if not venv_pip_version:
+            return Error(
+                "Failed to install pip into venv at {venv_dir}".format(venv_dir=venv.venv_dir)
+            )
+
+    if InstallScope.SOURCE_ONLY == scope:
+        return venv_pip_version
+
+    uninstall = OrderedDict()
+    pex_pip_version = find_dist(_PIP, pex.resolve())
+    if pex_pip_version and pex_pip_version != venv_pip_version:
+        uninstall[_PIP] = pex_pip_version
+
+    venv_setuptools_version = find_dist(_SETUPTOOLS, venv.iter_distributions())
+    if venv_setuptools_version:
+        pex_setuptools_version = find_dist(_SETUPTOOLS, pex.resolve())
+        if pex_setuptools_version and venv_setuptools_version != pex_setuptools_version:
+            uninstall[_SETUPTOOLS] = pex_setuptools_version
+
+    if not uninstall:
+        return venv_pip_version
+
+    message = (
+        "You asked for --pip to be installed in the venv at {venv_dir},\n"
+        "but the PEX at {pex} already contains:\n{distributions}"
+    ).format(
+        venv_dir=venv.venv_dir,
+        pex=pex.path(),
+        distributions="\n".join(
+            "{project_name} {version}".format(project_name=project_name, version=version)
+            for project_name, version in uninstall.items()
+        ),
+    )
+    if not collisions_ok:
+        return Error(
+            "{message}\nConsider re-running either without --pip or with --collisions-ok.".format(
+                message=message
+            )
+        )
+
+    pex_warnings.warn(
+        "{message}\nUninstalling venv versions and using versions from the PEX.".format(
+            message=message
+        )
+    )
+    projects_to_uninstall = sorted(str(project_name) for project_name in uninstall)
+    try:
+        subprocess.check_call(
+            args=[venv.interpreter.binary, "-m", "pip", "uninstall", "-y"] + projects_to_uninstall
+        )
+    except CalledProcessError as e:
+        return Error(
+            "Failed to uninstall venv versions of {projects}: {err}".format(
+                projects=" and ".join(projects_to_uninstall), err=e
+            )
+        )
+    return pex_pip_version or venv_pip_version
+
+
 class Venv(PEXCommand):
     """Creates a venv from the PEX file."""
 
@@ -133,7 +238,12 @@ def add_arguments(cls, parser):
             "--pip",
             action="store_true",
             default=False,
-            help="Add pip to the venv.",
+            help=(
+                "Add pip (and setuptools) to the venv. If the PEX already contains its own "
+                "conflicting versions pip (or setuptools), the command will error and you must "
+                "pass --collisions-ok to have the PEX versions over-ride the natural venv versions "
+                "installed by --pip."
+            ),
         )
         parser.add_argument(
             "--copies",
@@ -195,6 +305,13 @@ def run(self, pex):
                 prompt=self.options.prompt,
             )
 
+        if self.options.pip:
+            try_(
+                ensure_pip_installed(
+                    venv, pex, scope=self.options.scope, collisions_ok=self.options.collisions_ok
+                )
+            )
+
         if self.options.prompt != venv.custom_prompt:
             logger.warning(
                 "Unable to apply custom --prompt {prompt!r} in {python} venv; continuing with the "
@@ -211,14 +328,7 @@ def run(self, pex):
             scope=self.options.scope,
             hermetic_scripts=self.options.hermetic_scripts,
         )
-        if self.options.pip:
-            try:
-                venv.install_pip()
-            except PipUnavailableError as e:
-                return Error(
-                    "The virtual environment was successfully created, but Pip was not "
-                    "installed:\n{}".format(e)
-                )
+
         if self.options.compile:
             try:
                 pex.interpreter.execute(["-m", "compileall", venv_dir])
diff --git a/pex/venv/virtualenv.py b/pex/venv/virtualenv.py
--- a/pex/venv/virtualenv.py
+++ b/pex/venv/virtualenv.py
@@ -375,5 +375,5 @@ def install_pip(self):
                         os.path.join(atomic_dir.work_dir, os.path.basename(get_pip)), "wb"
                     ) as dst_fp:
                         shutil.copyfileobj(src_fp, dst_fp)
-            self._interpreter.execute(args=[get_pip])
+            self._interpreter.execute(args=[get_pip, "--no-wheel"])
         return self.bin_path("pip")


Test patch:
diff --git a/tests/integration/tools/commands/test_issue_2105.py b/tests/integration/tools/commands/test_issue_2105.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/tools/commands/test_issue_2105.py
@@ -0,0 +1,294 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import os.path
+import subprocess
+from textwrap import dedent
+
+import pytest
+
+from pex.dist_metadata import Distribution
+from pex.pep_440 import Version
+from pex.pep_503 import ProjectName
+from pex.pex import PEX
+from pex.testing import make_env, run_pex_command
+from pex.typing import TYPE_CHECKING
+from pex.venv.virtualenv import Virtualenv
+
+if TYPE_CHECKING:
+    from typing import Any, Iterable, Mapping
+
+
+@pytest.fixture(scope="module")
+def td(tmpdir_factory):
+    # type: (Any) -> Any
+    return tmpdir_factory.mktemp("td")
+
+
+PIP_PROJECT_NAME = ProjectName("pip")
+SETUPTOOLS_PROJECT_NAME = ProjectName("setuptools")
+
+
+def index_distributions(dists):
+    # type: (Iterable[Distribution]) -> Mapping[ProjectName, Version]
+    return {dist.metadata.project_name: dist.metadata.version for dist in dists}
+
+
+@pytest.fixture(scope="module")
+def baseline_venv_with_pip(td):
+    # type: (Any) -> Mapping[ProjectName, Version]
+    baseline_venv = Virtualenv.create(venv_dir=str(td.join("baseline.venv")))
+    baseline_venv.install_pip()
+    baseline_venv_distributions = index_distributions(baseline_venv.iter_distributions())
+    assert {PIP_PROJECT_NAME, SETUPTOOLS_PROJECT_NAME} == set(baseline_venv_distributions)
+    return baseline_venv_distributions
+
+
+@pytest.fixture(scope="module")
+def baseline_venv_pip_version(baseline_venv_with_pip):
+    # type: (Mapping[ProjectName, Version]) -> Version
+    return baseline_venv_with_pip[PIP_PROJECT_NAME]
+
+
+@pytest.fixture(scope="module")
+def baseline_venv_setuptools_version(baseline_venv_with_pip):
+    # type: (Mapping[ProjectName, Version]) -> Version
+    return baseline_venv_with_pip[SETUPTOOLS_PROJECT_NAME]
+
+
+def assert_venv_dists(
+    venv_dir,  # type: str
+    expected_pip_version,  # type: Version
+    expected_setuptools_version,  # type: Version
+):
+    virtualenv = Virtualenv(venv_dir)
+    dists = index_distributions(virtualenv.iter_distributions())
+    assert expected_pip_version == dists[PIP_PROJECT_NAME]
+    assert expected_setuptools_version == dists[SETUPTOOLS_PROJECT_NAME]
+
+    def reported_version(module):
+        # type: (str) -> Version
+        return Version(
+            subprocess.check_output(
+                args=[
+                    virtualenv.interpreter.binary,
+                    "-c",
+                    "import {module}; print({module}.__version__)".format(module=module),
+                ]
+            ).decode("utf-8")
+        )
+
+    assert expected_pip_version == reported_version("pip")
+    assert expected_setuptools_version == reported_version("setuptools")
+
+
+def assert_venv_dists_no_conflicts(
+    tmpdir,  # type: Any
+    pex,  # type: str
+    expected_pip_version,  # type: Version
+    expected_setuptools_version,  # type: Version
+):
+    # type: (...) -> None
+    venv_dir = os.path.join(str(tmpdir), "venv_dir")
+    subprocess.check_call(args=[pex, "venv", "--pip", venv_dir], env=make_env(PEX_TOOLS=1))
+    assert_venv_dists(venv_dir, expected_pip_version, expected_setuptools_version)
+
+
+def test_pip_empty_pex(
+    tmpdir,  # type: Any
+    baseline_venv_pip_version,  # type: Version
+    baseline_venv_setuptools_version,  # type: Version
+):
+    # type: (...) -> None
+
+    pex = os.path.join(str(tmpdir), "pex")
+    run_pex_command(args=["-o", pex, "--include-tools"]).assert_success()
+
+    assert_venv_dists_no_conflicts(
+        tmpdir,
+        pex,
+        expected_pip_version=baseline_venv_pip_version,
+        expected_setuptools_version=baseline_venv_setuptools_version,
+    )
+
+
+def test_pip_pex_no_conflicts(
+    tmpdir,  # type: Any
+    baseline_venv_pip_version,  # type: Version
+    baseline_venv_setuptools_version,  # type: Version
+):
+    # type: (...) -> None
+
+    pex = os.path.join(str(tmpdir), "pex")
+    run_pex_command(
+        args=[
+            "-o",
+            pex,
+            "pip=={version}".format(version=baseline_venv_pip_version),
+            "setuptools=={version}".format(version=baseline_venv_setuptools_version),
+            "--include-tools",
+        ]
+    ).assert_success()
+
+    assert_venv_dists_no_conflicts(
+        tmpdir,
+        pex,
+        expected_pip_version=baseline_venv_pip_version,
+        expected_setuptools_version=baseline_venv_setuptools_version,
+    )
+
+
+def assert_venv_dists_conflicts(
+    tmpdir,  # type: Any
+    pex,  # type: str
+    baseline_venv_pip_version,  # type: Version
+    baseline_venv_setuptools_version,  # type: Version
+    expected_pip_version,  # type: Version
+    expected_setuptools_version,  # type: Version
+):
+    # type: (...) -> None
+
+    expected_conflicts = []
+    if baseline_venv_pip_version != expected_pip_version:
+        expected_conflicts.append("pip {version}".format(version=expected_pip_version))
+    if baseline_venv_setuptools_version != expected_setuptools_version:
+        expected_conflicts.append(
+            "setuptools {version}".format(version=expected_setuptools_version)
+        )
+    assert (
+        expected_conflicts
+    ), "The assert_venv_dists_conflicts function requires at least one conflict."
+
+    venv_dir = os.path.join(str(tmpdir), "venv_dir")
+    args = [pex, "venv", "--pip", venv_dir]
+
+    expected_message_prefix = (
+        dedent(
+            """\
+            You asked for --pip to be installed in the venv at {venv_dir},
+            but the PEX at {pex} already contains:
+            {conflicts}
+            """
+        )
+        .format(venv_dir=venv_dir, pex=pex, conflicts=os.linesep.join(expected_conflicts))
+        .strip()
+    )
+
+    process = subprocess.Popen(args, stderr=subprocess.PIPE, env=make_env(PEX_TOOLS=1))
+    _, stderr = process.communicate()
+    assert 0 != process.returncode
+
+    decoded_stderr = stderr.decode("utf-8")
+    assert (
+        dedent(
+            """\
+            {prefix}
+            Consider re-running either without --pip or with --collisions-ok.
+            """
+        ).format(prefix=expected_message_prefix)
+        in decoded_stderr
+    ), decoded_stderr
+
+    process = subprocess.Popen(
+        args + ["--force", "--collisions-ok"], stderr=subprocess.PIPE, env=make_env(PEX_TOOLS=1)
+    )
+    _, stderr = process.communicate()
+    assert 0 == process.returncode
+    decoded_stderr = stderr.decode("utf-8")
+    assert (
+        dedent(
+            """\
+            {prefix}
+            Uninstalling venv versions and using versions from the PEX.
+            """
+        ).format(prefix=expected_message_prefix)
+        in decoded_stderr
+    ), decoded_stderr
+
+    assert_venv_dists(venv_dir, expected_pip_version, expected_setuptools_version)
+
+
+def test_pip_pex_pip_conflict(
+    tmpdir,  # type: Any
+    baseline_venv_pip_version,  # type: Version
+    baseline_venv_setuptools_version,  # type: Version
+):
+    # type: (...) -> None
+
+    pex = os.path.join(str(tmpdir), "pex")
+    run_pex_command(
+        args=[
+            "-o",
+            pex,
+            "pip!={version}".format(version=baseline_venv_pip_version),
+            "--include-tools",
+        ]
+    ).assert_success()
+    pex_pip_version = index_distributions(PEX(pex).resolve())[PIP_PROJECT_NAME]
+
+    assert_venv_dists_conflicts(
+        tmpdir,
+        pex,
+        baseline_venv_pip_version=baseline_venv_pip_version,
+        baseline_venv_setuptools_version=baseline_venv_setuptools_version,
+        expected_pip_version=pex_pip_version,
+        expected_setuptools_version=baseline_venv_setuptools_version,
+    )
+
+
+def test_pip_pex_setuptools_conflict(
+    tmpdir,  # type: Any
+    baseline_venv_pip_version,  # type: Version
+    baseline_venv_setuptools_version,  # type: Version
+):
+    # type: (...) -> None
+
+    pex = os.path.join(str(tmpdir), "pex")
+    run_pex_command(
+        args=[
+            "-o",
+            pex,
+            "setuptools!={version}".format(version=baseline_venv_setuptools_version),
+            "--include-tools",
+        ]
+    ).assert_success()
+    pex_setuptools_version = index_distributions(PEX(pex).resolve())[SETUPTOOLS_PROJECT_NAME]
+
+    assert_venv_dists_conflicts(
+        tmpdir,
+        pex,
+        baseline_venv_pip_version=baseline_venv_pip_version,
+        baseline_venv_setuptools_version=baseline_venv_setuptools_version,
+        expected_pip_version=baseline_venv_pip_version,
+        expected_setuptools_version=pex_setuptools_version,
+    )
+
+
+def test_pip_pex_both_conflict(
+    tmpdir,  # type: Any
+    baseline_venv_pip_version,  # type: Version
+    baseline_venv_setuptools_version,  # type: Version
+):
+    # type: (...) -> None
+
+    pex = os.path.join(str(tmpdir), "pex")
+    run_pex_command(
+        args=[
+            "-o",
+            pex,
+            "pip!={version}".format(version=baseline_venv_pip_version),
+            "setuptools!={version}".format(version=baseline_venv_setuptools_version),
+            "--include-tools",
+        ]
+    ).assert_success()
+    pex_pip_version = index_distributions(PEX(pex).resolve())[PIP_PROJECT_NAME]
+    pex_setuptools_version = index_distributions(PEX(pex).resolve())[SETUPTOOLS_PROJECT_NAME]
+
+    assert_venv_dists_conflicts(
+        tmpdir,
+        pex,
+        baseline_venv_pip_version=baseline_venv_pip_version,
+        baseline_venv_setuptools_version=baseline_venv_setuptools_version,
+        expected_pip_version=pex_pip_version,
+        expected_setuptools_version=pex_setuptools_version,
+    )