Implement support for preserving and injecting Python args.

pex/bin/pex.py

@@ -512,6 +512,18 @@ class InjectArgAction(Action):
         def __call__(self, parser, namespace, value, option_str=None):
             self.default.extend(shlex.split(value))
 
+    group.add_argument(
+        "--inject-python-args",
+        dest="inject_python_args",
+        default=[],
+        action=InjectArgAction,
+        help=(
+            "Command line arguments to the Python interpreter to freeze in. For example, `-u` to "
+            "disable buffering of `sys.stdout` and `sys.stderr` or `-W <arg>` to control Python "
+            "warnings."
+        ),
+    )
+
     group.add_argument(
         "--inject-args",
         dest="inject_args",
@@ -867,6 +879,7 @@ def build_pex(
             seen.add((src, dst))
 
     pex_info = pex_builder.info
+    pex_info.inject_python_args = options.inject_python_args
     pex_info.inject_env = dict(options.inject_env)
     pex_info.inject_args = options.inject_args
     pex_info.venv = bool(options.venv)

pex/build_system/pep_517.py

@@ -150,31 +150,33 @@ def _invoke_build_hook(
     build_backend_object = build_system.build_backend.replace(":", ".")
     with named_temporary_file(mode="r") as fp:
         args = build_system.venv_pex.execute_args(
-            "-c",
-            dedent(
-                """\
-                import json
-                import sys
+            additional_args=(
+                "-c",
+                dedent(
+                    """\
+                    import json
+                    import sys
 
-                import {build_backend_module}
+                    import {build_backend_module}
 
 
-                if not hasattr({build_backend_object}, {hook_method!r}):
-                    sys.exit({hook_unavailable_exit_code})
+                    if not hasattr({build_backend_object}, {hook_method!r}):
+                        sys.exit({hook_unavailable_exit_code})
 
-                result = {build_backend_object}.{hook_method}(*{hook_args!r}, **{hook_kwargs!r})
-                with open({result_file!r}, "w") as fp:
-                    json.dump(result, fp)
-                """
-            ).format(
-                build_backend_module=build_backend_module,
-                build_backend_object=build_backend_object,
-                hook_method=hook_method,
-                hook_args=tuple(hook_args),
-                hook_kwargs=dict(hook_kwargs) if hook_kwargs else {},
-                hook_unavailable_exit_code=_HOOK_UNAVAILABLE_EXIT_CODE,
-                result_file=fp.name,
-            ),
+                    result = {build_backend_object}.{hook_method}(*{hook_args!r}, **{hook_kwargs!r})
+                    with open({result_file!r}, "w") as fp:
+                        json.dump(result, fp)
+                    """
+                ).format(
+                    build_backend_module=build_backend_module,
+                    build_backend_object=build_backend_object,
+                    hook_method=hook_method,
+                    hook_args=tuple(hook_args),
+                    hook_kwargs=dict(hook_kwargs) if hook_kwargs else {},
+                    hook_unavailable_exit_code=_HOOK_UNAVAILABLE_EXIT_CODE,
+                    result_file=fp.name,
+                ),
+            )
         )
         process = subprocess.Popen(
             args=args,

pex/pex_boot.py

@@ -0,0 +1,219 @@
+# Copyright 2014 Pex project contributors.
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import os
+import sys
+
+TYPE_CHECKING = False
+
+if TYPE_CHECKING:
+    from typing import List, NoReturn, Optional, Tuple
+
+
+if sys.version_info >= (3, 10):
+
+    def orig_argv():
+        # type: () -> List[str]
+        return sys.orig_argv
+
+else:
+    try:
+        import ctypes
+
+        # N.B.: None of the PyPy versions we support <3.10 supports the pythonapi.
+        from ctypes import pythonapi
+
+        def orig_argv():
+            # type: () -> List[str]
+
+            # Under MyPy for Python 3.5, ctypes.POINTER is incorrectly typed. This code is tested
+            # to work correctly in practice on all Pythons Pex supports.
+            argv = ctypes.POINTER(  # type: ignore[call-arg]
+                ctypes.c_char_p if sys.version_info[0] == 2 else ctypes.c_wchar_p
+            )()
+
+            argc = ctypes.c_int()
+            pythonapi.Py_GetArgcArgv(ctypes.byref(argc), ctypes.byref(argv))
+
+            # Under MyPy for Python 3.5, argv[i] has its type incorrectly evaluated. This code
+            # is tested to work correctly in practice on all Pythons Pex supports.
+            return [argv[i] for i in range(argc.value)]  # type: ignore[misc]
+
+    except ImportError:
+        # N.B.: This handles the older PyPy case.
+        def orig_argv():
+            # type: () -> List[str]
+            return []
+
+
+def __re_exec__(
+    python,  # type: str
+    python_args,  # type: List[str]
+    *extra_python_args  # type: str
+):
+    # type: (...) -> NoReturn
+
+    argv = [python]
+    argv.extend(python_args)
+    argv.extend(extra_python_args)
+
+    os.execv(python, argv + sys.argv[1:])
+
+
+__SHOULD_EXECUTE__ = __name__ == "__main__"
+
+
+def __entry_point_from_filename__(filename):
+    # type: (str) -> str
+
+    # Either the entry point is "__main__" and we're in execute mode or "__pex__/__init__.py"
+    # and we're in import hook mode.
+    entry_point = os.path.dirname(filename)
+    if __SHOULD_EXECUTE__:
+        return entry_point
+    return os.path.dirname(entry_point)
+
+
+__INSTALLED_FROM__ = "__PEX_EXE__"
+
+
+def __ensure_pex_installed__(
+    pex,  # type: str
+    pex_root,  # type: str
+    pex_hash,  # type: str
+    python_args,  # type: List[str]
+):
+    # type: (...) -> Optional[str]
+
+    from pex.layout import ensure_installed
+    from pex.tracer import TRACER
+
+    installed_location = ensure_installed(pex=pex, pex_root=pex_root, pex_hash=pex_hash)
+    if not __SHOULD_EXECUTE__ or pex == installed_location:
+        return installed_location
+
+    # N.B.: This is read upon re-exec below to point sys.argv[0] back to the original pex
+    # before unconditionally scrubbing the env var and handing off to user code.
+    os.environ[__INSTALLED_FROM__] = pex
+
+    TRACER.log(
+        "Executing installed PEX for {pex} at {installed_location}".format(
+            pex=pex, installed_location=installed_location
+        )
+    )
+    __re_exec__(sys.executable, python_args, installed_location)
+
+
+def __maybe_run_venv__(
+    pex,  # type: str
+    pex_root,  # type: str
+    pex_hash,  # type: str
+    has_interpreter_constraints,  # type: bool
+    hermetic_venv_scripts,  # type: bool
+    pex_path,  # type: Tuple[str, ...]
+    python_args,  # type: List[str]
+):
+    # type: (...) -> Optional[str]
+
+    from pex.common import is_exe
+    from pex.tracer import TRACER
+    from pex.variables import venv_dir
+
+    venv_root_dir = venv_dir(
+        pex_file=pex,
+        pex_root=pex_root,
+        pex_hash=pex_hash,
+        has_interpreter_constraints=has_interpreter_constraints,
+        pex_path=pex_path,
+    )
+    venv_pex = os.path.join(venv_root_dir, "pex")
+    if not __SHOULD_EXECUTE__ or not is_exe(venv_pex):
+        # Code in bootstrap_pex will (re)create the venv after selecting the correct
+        # interpreter.
+        return venv_root_dir
+
+    TRACER.log("Executing venv PEX for {pex} at {venv_pex}".format(pex=pex, venv_pex=venv_pex))
+    venv_python = os.path.join(venv_root_dir, "bin", "python")
+    if hermetic_venv_scripts:
+        __re_exec__(venv_python, python_args, "-sE", venv_pex)
+    else:
+        __re_exec__(venv_python, python_args, venv_pex)
+
+
+def boot(
+    bootstrap_dir,  # type: str
+    pex_root,  # type: str
+    pex_hash,  # type: str
+    has_interpreter_constraints,  # type: bool
+    hermetic_venv_scripts,  # type: bool
+    pex_path,  # type: Tuple[str, ...]
+    is_venv,  # type: bool
+    inject_python_args,  # type: Tuple[str, ...]
+):
+    # type: (...) -> int
+
+    entry_point = None  # type: Optional[str]
+    __file__ = globals().get("__file__")
+    __loader__ = globals().get("__loader__")
+    if __file__ is not None and os.path.exists(__file__):
+        entry_point = __entry_point_from_filename__(__file__)
+    elif __loader__ is not None:
+        if hasattr(__loader__, "archive"):
+            entry_point = __loader__.archive
+        elif hasattr(__loader__, "get_filename"):
+            # The source of the loader interface has changed over the course of Python history
+            # from `pkgutil.ImpLoader` to `importlib.abc.Loader`, but the existence and
+            # semantics of `get_filename` has remained constant; so we just check for the
+            # method.
+            entry_point = __entry_point_from_filename__(__loader__.get_filename())
+
+    if entry_point is None:
+        sys.stderr.write("Could not launch python executable!\\n")
+        return 2
+
+    python_args = list(inject_python_args)  # type: List[str]
+    orig_args = orig_argv()
+    if orig_args:
+        for index, arg in enumerate(orig_args[1:], start=1):
+            if os.path.exists(arg) and os.path.samefile(entry_point, arg):
+                python_args.extend(orig_args[1:index])
+                break
+
+    installed_from = os.environ.pop(__INSTALLED_FROM__, None)
+    sys.argv[0] = installed_from or sys.argv[0]
+
+    sys.path[0] = os.path.abspath(sys.path[0])
+    sys.path.insert(0, os.path.abspath(os.path.join(entry_point, bootstrap_dir)))
+
+    venv_dir = None  # type: Optional[str]
+    if not installed_from:
+        os.environ["PEX"] = os.path.realpath(entry_point)
+        from pex.variables import ENV, Variables
+
+        pex_root = Variables.PEX_ROOT.value_or(ENV, pex_root)
+
+        if not ENV.PEX_TOOLS and Variables.PEX_VENV.value_or(ENV, is_venv):
+            venv_dir = __maybe_run_venv__(
+                pex=entry_point,
+                pex_root=pex_root,
+                pex_hash=pex_hash,
+                has_interpreter_constraints=has_interpreter_constraints,
+                hermetic_venv_scripts=hermetic_venv_scripts,
+                pex_path=ENV.PEX_PATH or pex_path,
+                python_args=python_args,
+            )
+        entry_point = __ensure_pex_installed__(
+            pex=entry_point, pex_root=pex_root, pex_hash=pex_hash, python_args=python_args
+        )
+        if entry_point is None:
+            # This means we re-exec'd ourselves already; so this just appeases type checking.
+            return 0
+    else:
+        os.environ["PEX"] = os.path.realpath(installed_from)
+
+    from pex.pex_bootstrapper import bootstrap_pex
+
+    bootstrap_pex(
+        entry_point, python_args=python_args, execute=__SHOULD_EXECUTE__, venv_dir=venv_dir
+    )
+    return 0

pex/pex_bootstrapper.py

@@ -27,7 +27,7 @@
 from pex.venv import installer
 
 if TYPE_CHECKING:
-    from typing import Iterable, Iterator, List, NoReturn, Optional, Set, Tuple, Union
+    from typing import Iterable, Iterator, List, NoReturn, Optional, Sequence, Set, Tuple, Union
 
     import attr  # vendor:skip
 
@@ -344,8 +344,11 @@ def gather_constraints():
         return target
 
 
-def maybe_reexec_pex(interpreter_test):
-    # type: (InterpreterTest) -> Union[None, NoReturn]
+def maybe_reexec_pex(
+    interpreter_test,  # type: InterpreterTest
+    python_args=(),  # type: Sequence[str]
+):
+    # type: (...) -> Union[None, NoReturn]
     """Handle environment overrides for the Python interpreter to use when executing this pex.
 
     This function supports interpreter filtering based on interpreter constraints stored in PEX-INFO
@@ -360,6 +363,7 @@ def maybe_reexec_pex(interpreter_test):
     constraints against these interpreters.
 
     :param interpreter_test: Optional test to verify selected interpreters can boot a given PEX.
+    :param python_args: Any args to pass to python when re-execing.
     """
 
     current_interpreter = PythonInterpreter.get()
@@ -411,7 +415,7 @@ def maybe_reexec_pex(interpreter_test):
         return None
 
     target_binary = target.binary
-    cmdline = [target_binary] + sys.argv
+    cmdline = [target_binary] + list(python_args) + sys.argv
     TRACER.log(
         "Re-executing: "
         "cmdline={cmdline!r}, "
@@ -461,18 +465,29 @@ def __attrs_post_init__(self):
         object.__setattr__(self, "pex", os.path.join(self.venv_dir, "pex"))
         object.__setattr__(self, "python", self.bin_file("python"))
 
-    def execute_args(self, *additional_args):
-        # type: (*str) -> List[str]
+    def execute_args(
+        self,
+        python_args=(),  # type: Sequence[str]
+        additional_args=(),  # type: Sequence[str]
+    ):
+        # type: (...) -> List[str]
         argv = [self.python]
+        argv.extend(python_args)
         if self.hermetic_scripts:
             argv.append("-sE")
         argv.append(self.pex)
         argv.extend(additional_args)
         return argv
 
-    def execv(self, *additional_args):
-        # type: (*str) -> NoReturn
-        os.execv(self.python, self.execute_args(*additional_args))
+    def execv(
+        self,
+        python_args=(),  # type: Sequence[str]
+        additional_args=(),  # type: Sequence[str]
+    ):
+        # type: (...) -> NoReturn
+        os.execv(
+            self.python, self.execute_args(python_args=python_args, additional_args=additional_args)
+        )
 
 
 def ensure_venv(
@@ -586,6 +601,7 @@ def bootstrap_pex(
     entry_point,  # type: str
     execute=True,  # type: bool
     venv_dir=None,  # type: Optional[str]
+    python_args=(),  # type: Sequence[str]
 ):
     # type: (...) -> None
 
@@ -619,9 +635,9 @@ def bootstrap_pex(
             except UnsatisfiableInterpreterConstraintsError as e:
                 die(str(e))
             venv_pex = _bootstrap_venv(entry_point, interpreter=target)
-            venv_pex.execv(*sys.argv[1:])
+            venv_pex.execv(python_args=python_args, additional_args=sys.argv[1:])
         else:
-            maybe_reexec_pex(interpreter_test=interpreter_test)
+            maybe_reexec_pex(interpreter_test=interpreter_test, python_args=python_args)
             from . import pex
 
             pex.PEX(entry_point).execute()