Fix PEXes for -i / PYTHONINSPECT=x. (#2491)

Previously PEXes did not behave like a Python interpreter when invoked
with either `-i` or with `PYTHONINSPECT=x`; now they do.

Fixes #2249

pex/globals.py

@@ -0,0 +1,6 @@
+# Copyright 2024 Pex project contributors.
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+
+class Globals(dict):
+    """The globals dict returned by PEX executions that evaluate code without exiting / exec'ing."""

pex/interpreter.py

@@ -938,6 +938,10 @@ def _create_isolated_cmd(
         pythonpath = list(pythonpath or ())
         if pythonpath:
             env["PYTHONPATH"] = os.pathsep.join(pythonpath)
+
+            # If we're being forced into interactive mode, we don't want that to apply to any
+            # Pex internal interpreter executions ever.
+            env.pop("PYTHONINSPECT", None)
         else:
             # Turn off reading of PYTHON* environment variables.
             cmd.append("-E")

pex/pex.py

@@ -19,6 +19,7 @@
 from pex.executor import Executor
 from pex.finders import get_entry_point_from_console_script, get_script_from_distributions
 from pex.fingerprinted_distribution import FingerprintedDistribution
+from pex.globals import Globals
 from pex.inherit_path import InheritPath
 from pex.interpreter import PythonIdentity, PythonInterpreter
 from pex.layout import Layout
@@ -536,7 +537,7 @@ def path(self):
         return self._pex
 
     def execute(self):
-        # type: () -> None
+        # type: () -> Any
         """Execute the PEX.
 
         This function makes assumptions that it is the last function called by the interpreter.
@@ -573,7 +574,11 @@ def execute(self):
                     V=3,
                 )
 
-        sys.exit(self._wrap_coverage(self._wrap_profiling, self._execute))
+        result = self._wrap_coverage(self._wrap_profiling, self._execute)
+        if "PYTHONINSPECT" not in os.environ:
+            sys.exit(0 if isinstance(result, Globals) else result)
+        else:
+            return result
 
     def _execute(self):
         # type: () -> Any
@@ -719,11 +724,12 @@ def execute_interpreter(self):
 
             import code
 
-            code.interact()
-            return None
+            local = {}  # type: Dict[str, Any]
+            code.interact(local=local)
+            return Globals(local)
 
+    @staticmethod
     def execute_with_options(
-        self,
         python_options,  # type: List[str]
         args,  # List[str]
     ):
@@ -744,6 +750,11 @@ def execute_with_options(
                 cmdline=" ".join(cmdline)
             )
         )
+        if any(
+            arg.startswith("-") and not arg.startswith("--") and "i" in arg
+            for arg in python_options
+        ):
+            os.environ["PYTHONINSPECT"] = "1"
         os.execv(python, cmdline)
 
     def execute_script(self, script_name):
@@ -786,7 +797,7 @@ def execute_content(
         content,  # type: str
         argv0=None,  # type: Optional[str]
     ):
-        # type: (...) -> Optional[str]
+        # type: (...) -> Any
         try:
             program = compile(content, name, "exec", flags=0, dont_inherit=1)
         except SyntaxError as e:
@@ -800,7 +811,7 @@ def execute_ast(
         program,  # type: ast.AST
         argv0=None,  # type: Optional[str]
     ):
-        # type: (...) -> Optional[str]
+        # type: (...) -> Any
         bootstrap.demote()
 
         from pex.compatibility import exec_function
@@ -809,8 +820,7 @@ def execute_ast(
         globals_map = globals().copy()
         globals_map["__name__"] = "__main__"
         globals_map["__file__"] = name
-        exec_function(program, globals_map)
-        return None
+        return Globals(exec_function(program, globals_map))
 
     def execute_entry(self, entry_point):
         # type: (Union[ModuleEntryPoint, CallableEntryPoint]) -> Any
@@ -820,12 +830,12 @@ def execute_entry(self, entry_point):
         return self.execute_module(entry_point.module)
 
     def execute_module(self, module_name):
-        # type: (str) -> None
+        # type: (str) -> Any
         bootstrap.demote()
 
         import runpy
 
-        runpy.run_module(module_name, run_name="__main__", alter_sys=True)
+        return Globals(runpy.run_module(module_name, run_name="__main__", alter_sys=True))
 
     @classmethod
     def execute_entry_point(cls, entry_point):

pex/pex_boot.py

@@ -7,7 +7,7 @@
 TYPE_CHECKING = False
 
 if TYPE_CHECKING:
-    from typing import List, NoReturn, Optional, Tuple
+    from typing import Any, List, NoReturn, Optional, Tuple
 
 
 if sys.version_info >= (3, 10):
@@ -150,7 +150,7 @@ def boot(
     is_venv,  # type: bool
     inject_python_args,  # type: Tuple[str, ...]
 ):
-    # type: (...) -> int
+    # type: (...) -> Tuple[Any, bool, bool]
 
     entry_point = None  # type: Optional[str]
     __file__ = globals().get("__file__")
@@ -169,7 +169,7 @@ def boot(
 
     if entry_point is None:
         sys.stderr.write("Could not launch python executable!\\n")
-        return 2
+        return 2, True, False
 
     python_args = list(inject_python_args)  # type: List[str]
     orig_args = orig_argv()
@@ -207,13 +207,16 @@ def boot(
         )
         if entry_point is None:
             # This means we re-exec'd ourselves already; so this just appeases type checking.
-            return 0
+            return 0, True, False
     else:
         os.environ["PEX"] = os.path.realpath(installed_from)
 
+    from pex.globals import Globals
     from pex.pex_bootstrapper import bootstrap_pex
 
-    bootstrap_pex(
+    result = bootstrap_pex(
         entry_point, python_args=python_args, execute=__SHOULD_EXECUTE__, venv_dir=venv_dir
     )
-    return 0
+    should_exit = __SHOULD_EXECUTE__ and "PYTHONINSPECT" not in os.environ
+    is_globals = isinstance(result, Globals)
+    return result, should_exit, is_globals

pex/pex_bootstrapper.py

@@ -27,7 +27,18 @@
 from pex.venv import installer
 
 if TYPE_CHECKING:
-    from typing import Iterable, Iterator, List, NoReturn, Optional, Sequence, Set, Tuple, Union
+    from typing import (
+        Any,
+        Iterable,
+        Iterator,
+        List,
+        NoReturn,
+        Optional,
+        Sequence,
+        Set,
+        Tuple,
+        Union,
+    )
 
     import attr  # vendor:skip
 
@@ -607,7 +618,7 @@ def bootstrap_pex(
     venv_dir=None,  # type: Optional[str]
     python_args=(),  # type: Sequence[str]
 ):
-    # type: (...) -> None
+    # type: (...) -> Any
 
     pex_info = _bootstrap(entry_point)
 
@@ -644,7 +655,7 @@ def bootstrap_pex(
             maybe_reexec_pex(interpreter_test=interpreter_test, python_args=python_args)
             from . import pex
 
-            pex.PEX(entry_point).execute()
+            return pex.PEX(entry_point).execute()
 
 
 def _activate_pex(

pex/pex_builder.py

@@ -486,7 +486,7 @@ def _prepare_code(self):
 
         pex_main = dedent(
             """
-            result = boot(
+            result, should_exit, is_globals = boot(
                 bootstrap_dir={bootstrap_dir!r},
                 pex_root={pex_root!r},
                 pex_hash={pex_hash!r},
@@ -496,8 +496,10 @@ def _prepare_code(self):
                 is_venv={is_venv!r},
                 inject_python_args={inject_python_args!r},
             )
-            if __SHOULD_EXECUTE__:
-                sys.exit(result)
+            if should_exit:
+                sys.exit(0 if is_globals else result)
+            elif is_globals:
+                globals().update(result)
             """
         ).format(
             bootstrap_dir=self._pex_info.bootstrap,

pex/scie/science.py

@@ -224,7 +224,7 @@ def create_cmd(named_entry_point):
                 "env": {
                     "default": env_default,
                     "remove_exact": ["PATH"],
-                    "remove_re": ["PEX_.*"],
+                    "remove_re": ["PEX_.*", "PYTHON.*"],
                     "replace": {
                         "PEX_INTERPRETER": "1",
                         # We can get a warning about too-long script shebangs, but this is not

pex/venv/installer.py

@@ -593,6 +593,9 @@ def _populate_first_party(
         """\
         {shebang}
 
+        from __future__ import print_function
+
+
         if __name__ == "__main__":
             import os
             import sys
@@ -629,12 +632,16 @@ def sys_executable_paths():
                         break
                 return executables
 
+            def maybe_log(*message):
+                if "PEX_VERBOSE" in os.environ:
+                    print(*message, file=sys.stderr)
+
             current_interpreter_blessed_env_var = "_PEX_SHOULD_EXIT_VENV_REEXEC"
             if (
                 not os.environ.pop(current_interpreter_blessed_env_var, None)
                 and sys_executable_paths().isdisjoint(iter_valid_venv_pythons())
             ):
-                sys.stderr.write("Re-execing from {{}}\\n".format(sys.executable))
+                maybe_log("Re-exec'ing from", sys.executable)
                 os.environ[current_interpreter_blessed_env_var] = "1"
                 argv = [python]
                 if {hermetic_re_exec!r}:
@@ -712,9 +719,9 @@ def sys_executable_paths():
                 )
             ]
             if ignored_pex_env_vars:
-                sys.stderr.write(
+                maybe_log(
                     "Ignoring the following environment variables in Pex venv mode:\\n"
-                    "{{}}\\n\\n".format(
+                    "{{}}\\n".format(
                         os.linesep.join(sorted(ignored_pex_env_vars))
                     )
                 )
@@ -859,21 +866,27 @@ def sys_executable_paths():
 
                 # The pex was called with Python interpreter options, so we need to re-exec to
                 # respect those:
-                if python_options:
-                    # Find the installed (unzipped) PEX entry point.
-                    main = sys.modules.get("__main__")
-                    if not main or not main.__file__:
-                        # N.B.: This should never happen.
-                        sys.stderr.write(
-                            "Unable to resolve PEX __main__ module file: {{}}\\n".format(main)
-                        )
-                        sys.exit(1)
-
+                if python_options or "PYTHONINSPECT" in os.environ:
                     python = sys.executable
-                    cmdline = [python] + python_options + [main.__file__] + args
-                    sys.stderr.write(
+                    cmdline = [python] + python_options
+                    inspect = "PYTHONINSPECT" in os.environ or any(
+                        arg.startswith("-") and not arg.startswith("--") and "i" in arg
+                        for arg in python_options
+                    )
+                    if not inspect:
+                        # We're not interactive; so find the installed (unzipped) PEX entry point.
+                        main = sys.modules.get("__main__")
+                        if not main or not main.__file__:
+                            # N.B.: This should never happen.
+                            sys.stderr.write(
+                                "Unable to resolve PEX __main__ module file: {{}}\\n".format(main)
+                            )
+                            sys.exit(1)
+                        cmdline.append(main.__file__)
+                    cmdline.extend(args)
+                    maybe_log(
                         "Re-executing with Python interpreter options: "
-                        "cmdline={{cmdline!r}}\\n".format(cmdline=" ".join(cmdline))
+                        "cmdline={{cmdline!r}}".format(cmdline=" ".join(cmdline))
                     )
                     os.execv(python, cmdline)
 

testing/scie.py

@@ -0,0 +1,21 @@
+# Copyright 2024 Pex project contributors.
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from pex.scie import SciePlatform
+from testing import IS_PYPY, PY_VER
+
+
+def has_provider():
+    # type: () -> bool
+    if IS_PYPY:
+        if PY_VER == (2, 7):
+            return True
+
+        if SciePlatform.LINUX_AARCH64 is SciePlatform.CURRENT:
+            return PY_VER >= (3, 7)
+        elif SciePlatform.MACOS_AARCH64 is SciePlatform.CURRENT:
+            return PY_VER >= (3, 8)
+        else:
+            return PY_VER >= (3, 6)
+    else:
+        return (3, 8) <= PY_VER < (3, 13)

tests/integration/scie/test_pex_scie.py

@@ -23,7 +23,7 @@
 from pex.targets import LocalInterpreter
 from pex.typing import TYPE_CHECKING
 from pex.version import __version__
-from testing import IS_PYPY, PY_VER, make_env, run_pex_command
+from testing import IS_PYPY, PY_VER, make_env, run_pex_command, scie
 
 if TYPE_CHECKING:
     from typing import Any, Iterable, List
@@ -488,24 +488,8 @@ def bar(tmpdir):
     return make_project(tmpdir, "bar")
 
 
-def has_provider():
-    # type: () -> bool
-    if IS_PYPY:
-        if PY_VER == (2, 7):
-            return True
-
-        if SciePlatform.LINUX_AARCH64 is SciePlatform.CURRENT:
-            return PY_VER >= (3, 7)
-        elif SciePlatform.MACOS_AARCH64 is SciePlatform.CURRENT:
-            return PY_VER >= (3, 8)
-        else:
-            return PY_VER >= (3, 6)
-    else:
-        return PY_VER >= (3, 8) and PY_VER < (3, 13)
-
-
 skip_if_no_provider = pytest.mark.skipif(
-    not has_provider(),
+    not scie.has_provider(),
     reason=(
         "Either A PBS or PyPy release must be available for the current interpreter to run this test."
     ),