Fix execution modes.

Previously you could not ask a `--venv` PEX to run in unzip mode. This
regularizes all PEX files so that they can be asked to operate in unzip
mode or venv mode whenever possible (venv mode requires the PEX was
built with `--include-tools` or `--venv`).

Work towards pex-tool#1343.

pex/bin/pex.py

@@ -947,7 +947,6 @@ def to_python_interpreter(full_path_or_basename):
         interpreter=interpreter,
         preamble=preamble,
         copy_mode=CopyMode.SYMLINK,
-        include_tools=options.include_tools or options.venv,
     )
 
     if options.resources_directory:
@@ -968,8 +967,9 @@ def to_python_interpreter(full_path_or_basename):
     pex_info.zip_safe = options.zip_safe
     pex_info.unzip = options.unzip
     pex_info.venv = bool(options.venv)
-    pex_info.venv_bin_path = options.venv
+    pex_info.venv_bin_path = options.venv or BinPath.FALSE
     pex_info.venv_copies = options.venv_copies
+    pex_info.includes_tools = options.include_tools or options.venv
     pex_info.pex_path = options.pex_path
     pex_info.always_write_cache = options.always_write_cache
     pex_info.ignore_errors = options.ignore_errors

pex/pex.py

@@ -469,15 +469,14 @@ def execute(self):
 
         try:
             if self._vars.PEX_TOOLS:
-                try:
-                    from pex.tools import main as tools
-                except ImportError as e:
+                if not self._pex_info.includes_tools:
                     die(
                         "The PEX_TOOLS environment variable was set, but this PEX was not built "
-                        "with tools (Re-build the PEX file with `pex --include-tools ...`):"
-                        " {}".format(e)
+                        "with tools (Re-build the PEX file with `pex --include-tools ...`)"
                     )
 
+                from pex.tools import main as tools
+
                 exit_value = tools.main(pex=self, pex_prog_path=sys.argv[0])
             else:
                 self.activate()

pex/pex_bootstrapper.py

@@ -369,6 +369,11 @@ def ensure_venv(pex):
         raise AssertionError(
             "Expected PEX-INFO for {} to have the components of a venv directory".format(pex.path())
         )
+    if not pex_info.includes_tools:
+        raise ValueError(
+            "The PEX_VENV environment variable was set, but this PEX was not built with venv "
+            "support (Re-build the PEX file with `pex --venv ...`)"
+        )
     with atomic_directory(venv_dir, exclusive=True) as venv:
         if not venv.is_finalized:
             from .tools.commands.venv import populate_venv_with_pex
@@ -456,7 +461,7 @@ def bootstrap_pex(entry_point):
     # ENV.PEX_ROOT is consulted by PythonInterpreter and Platform so set that up as early as
     # possible in the run.
     with ENV.patch(PEX_ROOT=pex_info.pex_root):
-        if not ENV.PEX_TOOLS and pex_info.venv:
+        if not (ENV.PEX_UNZIP or ENV.PEX_TOOLS) and pex_info.venv:
             try:
                 target = find_compatible_interpreter(
                     interpreter_constraints=pex_info.interpreter_constraints,
@@ -465,7 +470,10 @@ def bootstrap_pex(entry_point):
                 die(str(e))
             from . import pex
 
-            venv_pex = ensure_venv(pex.PEX(entry_point, interpreter=target))
+            try:
+                venv_pex = ensure_venv(pex.PEX(entry_point, interpreter=target))
+            except ValueError as e:
+                die(str(e))
             os.execv(venv_pex, [venv_pex] + sys.argv[1:])
         else:
             maybe_reexec_pex(pex_info.interpreter_constraints)

pex/pex_builder.py

@@ -69,12 +69,12 @@ def for_value(cls, value):
 import sys
 
 
-def __maybe_run_unzipped__(pex_zip):
+def __maybe_run_unzipped__(pex_zip, pex_root):
   from pex.common import atomic_directory, open_zip
   from pex.tracer import TRACER
   from pex.variables import unzip_dir
 
-  unzip_to = unzip_dir({pex_root!r}, {pex_hash!r})
+  unzip_to = unzip_dir(pex_root, {pex_hash!r})
   with atomic_directory(unzip_to, exclusive=True) as chroot:
     if not chroot.is_finalized:
       with TRACER.timed('Extracting {{}} to {{}}'.format(pex_zip, unzip_to)):
@@ -89,17 +89,17 @@ def __maybe_run_unzipped__(pex_zip):
   os.execv(sys.executable, [sys.executable, unzip_to] + sys.argv[1:])
 
 
-def __maybe_run_venv__(pex):
+def __maybe_run_venv__(pex, pex_root, pex_path):
   from pex.common import is_exe
   from pex.tracer import TRACER
   from pex.variables import venv_dir
 
   venv_home = venv_dir(
-    pex_root={pex_root!r}, 
+    pex_root=pex_root, 
     pex_hash={pex_hash!r},
     interpreter_constraints={interpreter_constraints!r},
     strip_pex_env={strip_pex_env!r},
-    pex_path={pex_path!r},
+    pex_path=pex_path,
   )
   venv_pex = os.path.join(venv_home, 'pex')
   if not is_exe(venv_pex):
@@ -133,19 +133,37 @@ def __maybe_run_venv__(pex):
 sys.path.insert(0, os.path.abspath(os.path.join(__entry_point__, {bootstrap_dir!r})))
 
 from pex.variables import ENV, Variables
-if Variables.PEX_VENV.value_or(ENV, {is_venv!r}):
-  if not {is_venv!r}:
-    from pex.common import die
-    die(
-      "The PEX_VENV environment variable was set, but this PEX was not built with venv support "
-      "(Re-build the PEX file with `pex --venv ...`):"
+__pex_root__ = Variables.PEX_ROOT.value_or(ENV, {pex_root!r})
+if ENV.PEX_VENV and ENV.PEX_UNZIP:
+  if {is_venv!r}:
+    default_mode = (
+      "Venv: Build a venv from the PEX zip file at {{}} under {{}} and run from there.".format(
+        __entry_point__, __pex_root__
+      )
     )
-  if not ENV.PEX_TOOLS:  # We need to run from the PEX for access to tools.
-    __maybe_run_venv__(__entry_point__)
+  elif {is_unzip!r}:
+    default_mode = "Unzip: Unzip the PEX zip file at {{}} under {{}} and run from there.".format(
+        __entry_point__, __pex_root__
+    )
+  else:
+    default_mode = "Zipapp: Run from the PEX zip file at {{}}.".format(__entry_point__)
+  from pex.common import die
+  die(
+    "Cannot request both PEX_VENV and PEX_UNZIP via environment variable.\\n"
+    "Please use just one of those environment variables or else unset both and accept this PEX's "
+    "default mode:\\n{{}}".format(default_mode)
+  )
+
+if not (ENV.PEX_UNZIP or ENV.PEX_TOOLS) and Variables.PEX_VENV.value_or(ENV, {is_venv!r}):
+  __maybe_run_venv__(
+    __entry_point__,
+    pex_root=__pex_root__,
+    pex_path=Variables.PEX_PATH.value_or(ENV, {pex_path!r}),
+  )
 elif Variables.PEX_UNZIP.value_or(ENV, {is_unzip!r}):
   import zipfile
   if zipfile.is_zipfile(__entry_point__):
-    __maybe_run_unzipped__(__entry_point__)
+    __maybe_run_unzipped__(__entry_point__, __pex_root__)
 
 from pex.pex_bootstrapper import bootstrap_pex
 bootstrap_pex(__entry_point__)
@@ -178,7 +196,6 @@ def __init__(
         pex_info=None,  # type: Optional[PexInfo]
         preamble=None,  # type: Optional[str]
         copy_mode=CopyMode.LINK,  # type: CopyMode.Value
-        include_tools=False,  # type: bool
     ):
         # type: (...) -> None
         """Initialize a pex builder.
@@ -192,8 +209,6 @@ def __init__(
         :keyword preamble: If supplied, execute this code prior to bootstrapping this PEX
           environment.
         :keyword copy_mode: Create the pex environment using the given copy mode.
-        :keyword include_tools: If True, include runtime tools which can be executed by exporting
-                                `PEX_TOOLS=1`.
 
         .. versionchanged:: 0.8
           The temporary directory created when ``path`` is not specified is now garbage collected on
@@ -204,7 +219,6 @@ def __init__(
         self._pex_info = pex_info or PexInfo.default(self._interpreter)
         self._preamble = preamble or ""
         self._copy_mode = copy_mode
-        self._include_tools = include_tools
 
         self._shebang = self._interpreter.identity.hashbang()
         self._logger = logging.getLogger(__name__)
@@ -568,7 +582,7 @@ def _prepare_bootstrap(self):
             # do use at runtime, does.
             root_module_names=["attr", "packaging", "pkg_resources", "pyparsing"],
         )
-        if self._include_tools:
+        if self._pex_info.includes_tools:
             # The `repository extract` tool needs setuptools and wheel to build sdists and wheels
             # and distutils needs .dist-info to discover setuptools (and wheel).
             vendor.vendor_runtime(
@@ -586,7 +600,7 @@ def _prepare_bootstrap(self):
             provider = ZipProvider(mod)
 
         bootstrap_packages = ["", "third_party"]
-        if self._include_tools:
+        if self._pex_info.includes_tools:
             bootstrap_packages.extend(["tools", "tools/commands"])
         for package in bootstrap_packages:
             for fn in provider.resource_listdir(package):

pex/pex_info.py

@@ -279,6 +279,16 @@ def venv_dir(self):
             pex_path=self.pex_path,
         )
 
+    @property
+    def includes_tools(self):
+        # type: () -> bool
+        return self._pex_info.get("includes_tools", self.venv)
+
+    @includes_tools.setter
+    def includes_tools(self, value):
+        # type: (bool) -> None
+        self._pex_info["includes_tools"] = bool(value)
+
     @property
     def strip_pex_env(self):
         """Whether or not this PEX should strip `PEX_*` env vars before executing its entrypoint.

pex/variables.py

@@ -523,7 +523,7 @@ def PEX_EXTRA_SYS_PATH(self):
         """
         return self._maybe_get_string("PEX_EXTRA_SYS_PATH")
 
-    @defaulted_property(default=os.path.expanduser("~/.pex"))
+    @defaulted_property(default="~/.pex")
     def PEX_ROOT(self):
         # type: () -> str
         """Directory.
@@ -536,7 +536,8 @@ def PEX_ROOT(self):
         return self._get_path("PEX_ROOT")
 
     @PEX_ROOT.validator
-    def _ensure_writeable_pex_root(self, pex_root):
+    def _ensure_writeable_pex_root(self, raw_pex_root):
+        pex_root = os.path.expanduser(raw_pex_root)
         if not can_write_dir(pex_root):
             tmp_root = os.path.realpath(safe_mkdtemp())
             pex_warnings.warn(

tests/test_execution_mode.py

@@ -0,0 +1,81 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import
+
+import os.path
+import subprocess
+from subprocess import CalledProcessError
+
+import pytest
+
+from pex.testing import run_pex_command
+from pex.typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Dict, Iterable, Tuple
+
+    CreateColorsPex = Callable[[Iterable[str]], str]
+    ExecuteColorsPex = Callable[[str, Dict[str, str]], Tuple[str, str]]
+
+
+@pytest.fixture
+def create_colors_pex(tmpdir):
+    # type: (Any) -> CreateColorsPex
+    def create(extra_args):
+        pex_file = os.path.join(str(tmpdir), "colors.pex")
+        results = run_pex_command(["ansicolors==1.1.8", "-o", pex_file] + list(extra_args))
+        results.assert_success()
+        return pex_file
+
+    return create
+
+
+@pytest.fixture
+def execute_colors_pex(tmpdir):
+    # type: (Any) -> ExecuteColorsPex
+    def execute(colors_pex, extra_env):
+        pex_root = os.path.join(str(tmpdir), "pex_root")
+        env = os.environ.copy()
+        env.update(extra_env)
+        env["PEX_ROOT"] = pex_root
+        output = subprocess.check_output(
+            [colors_pex, "-c", "import colors; print(colors.__file__)"], env=env
+        )
+        return output.strip().decode("utf-8"), pex_root
+
+    return execute
+
+
+@pytest.mark.parametrize(
+    ["extra_args", "default_dir", "venv_exception_expected"],
+    [
+        pytest.param([], "installed_wheels", True, id="ZIPAPP"),
+        pytest.param(["--include-tools"], "installed_wheels", False, id="ZIPAPP --include-tools"),
+        pytest.param(["--unzip"], "unzipped", True, id="UNZIP"),
+        pytest.param(["--unzip", "--include-tools"], "unzipped", False, id="UNZIP --include-tools"),
+        pytest.param(["--venv"], "venvs", False, id="VENV"),
+    ],
+)
+def test_execution_mode(
+    create_colors_pex,  # type: CreateColorsPex
+    execute_colors_pex,  # type: ExecuteColorsPex
+    extra_args,  # type: Iterable[str]
+    default_dir,  # type: str
+    venv_exception_expected,  # type: bool
+):
+    # type: (...) -> None
+    pex_file = create_colors_pex(extra_args)
+
+    output, pex_root = execute_colors_pex(pex_file, {})
+    assert output.startswith(os.path.join(pex_root, default_dir))
+
+    output, pex_root = execute_colors_pex(pex_file, {"PEX_UNZIP": "1"})
+    assert output.startswith(os.path.join(pex_root, "unzipped"))
+
+    if venv_exception_expected:
+        with pytest.raises(CalledProcessError):
+            execute_colors_pex(pex_file, {"PEX_VENV": "1"})
+    else:
+        output, pex_root = execute_colors_pex(pex_file, {"PEX_VENV": "1"})
+        assert output.startswith(os.path.join(pex_root, "venvs"))

tests/tools/commands/test_interpreter_command.py

@@ -47,7 +47,8 @@ def create(
         *other_interpreters  # type: PythonInterpreter
     ):
         # type: (...) -> InterpreterTool
-        pex_builder = PEXBuilder(include_tools=True, interpreter=interpreter)
+        pex_builder = PEXBuilder(interpreter=interpreter)
+        pex_builder.info.includes_tools = True
         pex_builder.freeze()
         return cls(
             tools_pex=pex_builder.path(),