Improve venv creation robustness when adding Pip. (#2454)

Fixes #2413

CHANGES.md

@@ -1,5 +1,27 @@
 # Release Notes
 
+## 2.8.0
+
+This release adds a new `--override` option to resolves that ultimately
+use an `--index` or `--find-links`. This allows you to override
+transitive dependencies when you have determined they are too narrow and
+that expanding their range is safe to do. The new `--override`s and the
+existing `--exclude`s can now also be specified when creating or syncing
+a lock file to seal these dependency modifications into the lock.
+
+This release also adds a new `--project` option to `pex` and
+`pex3 lock {create,sync}` that improves the ergonomics of locking a
+local Python project and then creating PEX executables for that project
+using its locked requirements.
+
+In addition, this release fixes the `bdist_pex` distutils command that
+ships with Pex to work when run under `tox` and Python 3.12 by improving
+Pex venv creation robustness when creating venvs that include Pip.
+
+* Add support for `--override`. (#2431)
+* Support `--project` locking and PEX building. (#2455)
+* Improve venv creation robustness when adding Pip. (#2454)
+
 ## 2.7.0
 
 This release adds support for Pip 24.1.1.

pex/build_system/pep_518.py

@@ -110,7 +110,7 @@ def create(
             # work well enough for our test cases and, in general, they should work well enough with
             # the Python they come paired with.
             upgrade_pip = virtualenv.interpreter.version[:2] == (3, 5)
-            virtualenv.install_pip(upgrade=upgrade_pip)
+            virtualenv.ensure_pip(upgrade=upgrade_pip)
             with open(os.devnull, "wb") as dev_null:
                 _, process = virtualenv.interpreter.open_process(
                     args=[

pex/cli/commands/lock.py

@@ -307,7 +307,7 @@ def sync(
                     os.rmdir(parent_dir)
 
         if retain_pip and not resolved_pip and not installed_pip:
-            self.venv.install_pip(upgrade=True)
+            self.venv.ensure_pip(upgrade=True)
 
         if to_install:
             for distribution in to_install:

pex/pip/installation.py

@@ -23,7 +23,7 @@
 from pex.typing import TYPE_CHECKING
 from pex.util import named_temporary_file
 from pex.variables import ENV
-from pex.venv.virtualenv import Virtualenv
+from pex.venv.virtualenv import InstallationChoice, Virtualenv
 
 if TYPE_CHECKING:
     from typing import Callable, Dict, Iterator, Optional, Union
@@ -99,8 +99,11 @@ def bootstrap_pip():
         # type: () -> Iterator[str]
 
         chroot = safe_mkdtemp()
-        venv = Virtualenv.create(venv_dir=os.path.join(chroot, "pip"), interpreter=interpreter)
-        venv.install_pip(upgrade=True)
+        venv = Virtualenv.create(
+            venv_dir=os.path.join(chroot, "pip"),
+            interpreter=interpreter,
+            install_pip=InstallationChoice.UPGRADED,
+        )
 
         for req in version.requirements:
             project_name = Requirement.parse(req).name

pex/tools/commands/repository.py

@@ -34,7 +34,7 @@
 from pex.tools.command import PEXCommand
 from pex.typing import TYPE_CHECKING, cast
 from pex.variables import ENV
-from pex.venv.virtualenv import Virtualenv
+from pex.venv.virtualenv import InstallationChoice, Virtualenv
 
 if TYPE_CHECKING:
     from typing import IO, Any, Callable, Iterable, Iterator, List, Text, Tuple
@@ -58,9 +58,13 @@ def spawn_python_job_with_setuptools_and_wheel(
     venv_dir = os.path.join(ENV.PEX_ROOT, "tools", "repository", str(interpreter.platform))
     with atomic_directory(venv_dir) as atomic_dir:
         if not atomic_dir.is_finalized():
-            venv = Virtualenv.create_atomic(venv_dir=atomic_dir, interpreter=interpreter)
-            venv.install_pip(upgrade=True)
-            venv.interpreter.execute(args=["-m", "pip", "install", "-U", "setuptools", "wheel"])
+            Virtualenv.create_atomic(
+                venv_dir=atomic_dir,
+                interpreter=interpreter,
+                install_pip=InstallationChoice.UPGRADED,
+                install_setuptools=InstallationChoice.UPGRADED,
+                install_wheel=InstallationChoice.UPGRADED,
+            )
 
     execute_python_args = [Virtualenv(venv_dir=venv_dir).interpreter.binary]
     execute_python_args.extend(args)

pex/venv/installer.py

@@ -79,7 +79,7 @@ def ensure_pip_installed(
         )
     else:
         try:
-            venv.install_pip()
+            venv.ensure_pip()
         except PipUnavailableError as e:
             return Error(
                 "The virtual environment was successfully created, but Pip was not "

pex/venv/virtualenv.py

@@ -18,6 +18,7 @@
 from pex.common import is_exe, safe_mkdir, safe_open
 from pex.compatibility import commonpath, get_stdout_bytes_buffer
 from pex.dist_metadata import Distribution, find_distributions
+from pex.enum import Enum
 from pex.executor import Executor
 from pex.fetcher import URLFetcher
 from pex.interpreter import (
@@ -138,6 +139,15 @@ def _find_preferred_site_packages_dir(
     )
 
 
+class InstallationChoice(Enum["InstallationChoice.Value"]):
+    class Value(Enum.Value):
+        pass
+
+    NO = Value("no")
+    YES = Value("yes")
+    UPGRADED = Value("upgraded")
+
+
 class Virtualenv(object):
     VIRTUALENV_VERSION = "16.7.12"
 
@@ -168,8 +178,29 @@ def create(
         copies=False,  # type: bool
         system_site_packages=False,  # type: bool
         prompt=None,  # type: Optional[str]
+        install_pip=InstallationChoice.NO,  # type: InstallationChoice.Value
+        install_setuptools=InstallationChoice.NO,  # type: InstallationChoice.Value
+        install_wheel=InstallationChoice.NO,  # type: InstallationChoice.Value
     ):
         # type: (...) -> Virtualenv
+
+        installations = {
+            "pip": install_pip,
+            "setuptools": install_setuptools,
+            "wheel": install_wheel,
+        }
+        project_upgrades = [
+            project
+            for project, installation_choice in installations.items()
+            if installation_choice is InstallationChoice.UPGRADED
+        ]
+        if project_upgrades and install_pip is InstallationChoice.NO:
+            raise ValueError(
+                "Installation of Pip is required in order to upgrade {projects}.".format(
+                    projects=" and ".join(project_upgrades)
+                )
+            )
+
         venv_dir = os.path.abspath(venv_dir)
         safe_mkdir(venv_dir, clean=force)
 
@@ -183,6 +214,31 @@ def create(
             )
             interpreter = base_interpreter
 
+        # N.B.: PyPy3.6 and PyPy3.7 come equipped with a venv module but it does not seem to
+        # work.
+        py_major_minor = interpreter.version[:2]
+        use_virtualenv = py_major_minor[0] == 2 or (
+            interpreter.is_pypy and py_major_minor[:2] <= (3, 7)
+        )
+
+        installations.pop("pip")
+        if install_pip is not InstallationChoice.NO and py_major_minor < (3, 12):
+            # The ensure_pip module, get_pip.py and the venv module all install setuptools when
+            # they install Pip for all Pythons older than 3.12.
+            installations.pop("setuptools")
+
+        project_installs = [
+            project
+            for project, installation_choice in installations.items()
+            if installation_choice is InstallationChoice.YES
+        ]
+        if project_installs and install_pip is InstallationChoice.NO:
+            raise ValueError(
+                "Installation of Pip is required in order to install {projects}.".format(
+                    projects=" and ".join(project_installs)
+                )
+            )
+
         # Guard against API calls from environment with ambient PYTHONPATH preventing pip virtualenv
         # creation. See: https://github.com/pex-tool/pex/issues/1451
         env = os.environ.copy()
@@ -196,10 +252,7 @@ def create(
             )
 
         custom_prompt = None  # type: Optional[str]
-        py_major_minor = interpreter.version[:2]
-        if py_major_minor[0] == 2 or (interpreter.is_pypy and py_major_minor[:2] <= (3, 7)):
-            # N.B.: PyPy3.6 and PyPy3.7 come equipped with a venv module but it does not seem to
-            # work.
+        if use_virtualenv:
             virtualenv_py = pkgutil.get_data(
                 __name__, "virtualenv_{version}_py".format(version=cls.VIRTUALENV_VERSION)
             )
@@ -243,7 +296,9 @@ def create(
                         )
                     )
         else:
-            args = ["-m", "venv", "--without-pip", venv_dir]
+            args = ["-m", "venv", venv_dir]
+            if install_pip is InstallationChoice.NO:
+                args.append("--without-pip")
             if copies:
                 args.append("--copies")
             if system_site_packages:
@@ -253,7 +308,20 @@ def create(
                 custom_prompt = prompt
             interpreter.execute(args=args, env=env)
 
-        return cls(venv_dir, custom_prompt=custom_prompt)
+        venv = cls(venv_dir, custom_prompt=custom_prompt)
+        if use_virtualenv and (
+            install_pip is not InstallationChoice.NO or project_upgrades or project_installs
+        ):
+            # Our vendored virtualenv does not support installing Pip, setuptool or wheel; so we
+            # use the ensurepip module / get_pip.py bootstrapping for Pip that `ensure_pip` does.
+            venv.ensure_pip(upgrade=install_pip is InstallationChoice.UPGRADED)
+        if project_upgrades:
+            venv.interpreter.execute(
+                args=["-m", "pip", "install", "-U"] + project_upgrades, env=env
+            )
+        if project_installs:
+            venv.interpreter.execute(args=["-m", "pip", "install"] + project_installs, env=env)
+        return venv
 
     @classmethod
     def create_atomic(
@@ -263,6 +331,9 @@ def create_atomic(
         force=False,  # type: bool
         copies=False,  # type: bool
         prompt=None,  # type: Optional[str]
+        install_pip=InstallationChoice.NO,  # type: InstallationChoice.Value
+        install_setuptools=InstallationChoice.NO,  # type: InstallationChoice.Value
+        install_wheel=InstallationChoice.NO,  # type: InstallationChoice.Value
     ):
         # type: (...) -> Virtualenv
         virtualenv = cls.create(
@@ -271,6 +342,9 @@ def create_atomic(
             force=force,
             copies=copies,
             prompt=prompt,
+            install_pip=install_pip,
+            install_setuptools=install_setuptools,
+            install_wheel=install_wheel,
         )
         for script in virtualenv._rewrite_base_scripts(real_venv_dir=venv_dir.target_dir):
             TRACER.log("Re-writing {}".format(script))
@@ -327,7 +401,6 @@ def __init__(
             elif isinstance(entry, Platlib):
                 self._platlib = entry.path
 
-        self._base_bin = frozenset(_iter_files(self._bin_dir))
         self._sys_path = None  # type: Optional[Tuple[str, ...]]
 
     @property
@@ -415,7 +488,7 @@ def _rewrite_base_scripts(self, real_venv_dir):
         # type: (str) -> Iterator[str]
         scripts = [
             path
-            for path in self._base_bin
+            for path in _iter_files(self._bin_dir)
             if _is_python_script(path) or re.search(r"^[Aa]ctivate", os.path.basename(path))
         ]
         if scripts:
@@ -456,8 +529,11 @@ def rewrite_scripts(
                         # N.B.: These lines include the newline already.
                         buffer.write(cast(bytes, line))
 
-    def install_pip(self, upgrade=False):
+    def ensure_pip(self, upgrade=False):
         # type: (bool) -> str
+        pip_script = self.bin_path("pip")
+        if is_exe(pip_script) and not upgrade:
+            return pip_script
         try:
             self._interpreter.execute(args=["-m", "ensurepip", "-U", "--default-pip"])
         except Executor.NonZeroExit:
@@ -487,4 +563,4 @@ def install_pip(self, upgrade=False):
             self._interpreter.execute(args=[get_pip, "--no-wheel"])
         if upgrade:
             self._interpreter.execute(args=["-m", "pip", "install", "-U", "pip"])
-        return self.bin_path("pip")
+        return pip_script

pex/version.py

@@ -1,4 +1,4 @@
 # Copyright 2015 Pex project contributors.
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-__version__ = "2.7.0"
+__version__ = "2.8.0"

testing/__init__.py

@@ -31,7 +31,7 @@
 from pex.resolve.resolver_configuration import PipConfiguration
 from pex.typing import TYPE_CHECKING
 from pex.util import named_temporary_file
-from pex.venv.virtualenv import Virtualenv
+from pex.venv.virtualenv import InstallationChoice, Virtualenv
 
 try:
     from unittest import mock
@@ -647,8 +647,8 @@ def python_venv(
         venv_dir=venv_dir or safe_mkdtemp(),
         interpreter=PythonInterpreter.from_binary(python),
         system_site_packages=system_site_packages,
+        install_pip=InstallationChoice.YES,
     )
-    venv.install_pip()
     return venv.interpreter.binary, venv.bin_path("pip")
 
 

tests/integration/cli/commands/test_interpreter_inspect.py

@@ -8,7 +8,7 @@
 from pex.pep_425 import CompatibilityTags
 from pex.pep_508 import MarkerEnvironment
 from pex.typing import TYPE_CHECKING, cast
-from pex.venv.virtualenv import Virtualenv
+from pex.venv.virtualenv import InstallationChoice, Virtualenv
 from testing import IntegResults
 from testing.cli import run_pex3
 
@@ -140,8 +140,9 @@ def test_inspect_interpreter_selection(
 def test_inspect_distributions(tmpdir):
     # type: (Any) -> None
 
-    venv = Virtualenv.create(venv_dir=os.path.join(str(tmpdir), "venv"))
-    venv.install_pip()
+    venv = Virtualenv.create(
+        venv_dir=os.path.join(str(tmpdir), "venv"), install_pip=InstallationChoice.YES
+    )
     venv.interpreter.execute(args=["-mpip", "install", "ansicolors==1.1.8", "cowsay==5.0"])
 
     data = assert_verbose_data(venv.interpreter, "-vd", "--python", venv.interpreter.binary)

tests/integration/cli/commands/test_local_project_lock.py

@@ -11,7 +11,7 @@
 from pex.common import touch
 from pex.interpreter import PythonInterpreter
 from pex.typing import TYPE_CHECKING
-from pex.venv.virtualenv import Virtualenv
+from pex.venv.virtualenv import InstallationChoice, Virtualenv
 from testing import PY27, ensure_python_interpreter, run_pex_command
 from testing.cli import run_pex3
 
@@ -57,9 +57,9 @@ def test_fingerprint_stability(
     tox_venv = Virtualenv.create(
         venv_dir=os.path.join(str(tmpdir), "tox.venv"),
         interpreter=PythonInterpreter.from_binary(ensure_python_interpreter(PY27)),
+        install_pip=InstallationChoice.YES,
     )
-    tox_venv.install_pip()
-    subprocess.check_call(args=[tox_venv.bin_path("pip"), "install", "tox"])
+    subprocess.check_call(args=[(tox_venv.bin_path("pip")), "install", "tox"])
     subprocess.check_call(args=[tox_venv.bin_path("tox"), "-e", "py27"], cwd=ansicolors_1_1_8)
     run_pex_command(args=print_colors_version_args).assert_success()
 

tests/integration/cli/commands/test_lock_sync.py

@@ -35,7 +35,7 @@
 from pex.sorted_tuple import SortedTuple
 from pex.third_party.packaging.tags import Tag
 from pex.typing import TYPE_CHECKING, cast
-from pex.venv.virtualenv import Virtualenv
+from pex.venv.virtualenv import InstallationChoice, Virtualenv
 from testing import (
     IS_PYPY,
     IS_X86_64,
@@ -998,8 +998,7 @@ def test_sync_venv_run_no_retain_pip_preinstalled(
     # type: (...) -> None
 
     venv_dir = os.path.join(str(tmpdir), "venv")
-    venv = Virtualenv.create(venv_dir)
-    venv.install_pip()
+    venv = Virtualenv.create(venv_dir, install_pip=InstallationChoice.YES)
     pip = find_distribution("pip", search_path=venv.sys_path)
     assert pip is not None
 
@@ -1037,8 +1036,7 @@ def test_sync_venv_run_retain_pip_preinstalled(
     # type: (...) -> None
 
     venv_dir = os.path.join(str(tmpdir), "venv")
-    venv = Virtualenv.create(venv_dir)
-    venv.install_pip()
+    venv = Virtualenv.create(venv_dir, install_pip=InstallationChoice.YES)
     pip = find_distribution("pip", search_path=venv.sys_path)
     assert pip is not None
     pip_pin = pin("pip", pip.version)
@@ -1149,8 +1147,7 @@ def test_sync_venv_run_retain_user_pip(
     # type: (...) -> None
 
     venv_dir = os.path.join(str(tmpdir), "venv")
-    venv = Virtualenv.create(venv_dir)
-    venv.install_pip()
+    venv = Virtualenv.create(venv_dir, install_pip=InstallationChoice.YES)
     original_pip = find_distribution("pip", search_path=venv.sys_path)
     assert original_pip is not None