diff --git a/pex/bin/pex.py b/pex/bin/pex.py index 1e41b5de0..c2b2c3a0c 100755 --- a/pex/bin/pex.py +++ b/pex/bin/pex.py @@ -513,6 +513,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 ` to control Python " + "warnings." + ), + ) + group.add_argument( "--inject-args", dest="inject_args", @@ -868,6 +880,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) diff --git a/pex/build_system/pep_517.py b/pex/build_system/pep_517.py index 837a5d35e..b085f47b2 100644 --- a/pex/build_system/pep_517.py +++ b/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, diff --git a/pex/pex_boot.py b/pex/pex_boot.py new file mode 100644 index 000000000..9e7148696 --- /dev/null +++ b/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 diff --git a/pex/pex_bootstrapper.py b/pex/pex_bootstrapper.py index c1ccebd51..a097736fa 100644 --- a/pex/pex_bootstrapper.py +++ b/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() diff --git a/pex/pex_builder.py b/pex/pex_builder.py index 7a9d5ee83..0a33d6e37 100644 --- a/pex/pex_builder.py +++ b/pex/pex_builder.py @@ -1,7 +1,7 @@ # Copyright 2014 Pex project contributors. # Licensed under the Apache License, Version 2.0 (see LICENSE). -from __future__ import absolute_import +from __future__ import absolute_import, print_function import hashlib import logging @@ -118,114 +118,6 @@ def perform_check( ERROR = Value("error") -BOOTSTRAP_ENVIRONMENT = """\ -import os -import sys - - -__INSTALLED_FROM__ = '__PEX_EXE__' - - -def __re_exec__(argv0, *extra_launch_args): - os.execv(argv0, [argv0] + list(extra_launch_args) + sys.argv[1:]) - - -__execute__ = __name__ == "__main__" - -def __ensure_pex_installed__(pex, pex_root, pex_hash): - from pex.layout import ensure_installed - from pex.tracer import TRACER - - installed_location = ensure_installed(pex, pex_root, pex_hash) - if not __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 {{}} at {{}}'.format(pex, installed_location)) - __re_exec__(sys.executable, installed_location) - - -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_dir = venv_dir( - pex_file=pex, - pex_root=pex_root, - pex_hash={pex_hash!r}, - has_interpreter_constraints={has_interpreter_constraints!r}, - pex_path=pex_path, - ) - venv_pex = os.path.join(venv_dir, 'pex') - if not __execute__ or not is_exe(venv_pex): - # Code in bootstrap_pex will (re)create the venv after selecting the correct interpreter. - return venv_dir - - TRACER.log('Executing venv PEX for {{}} at {{}}'.format(pex, venv_pex)) - venv_python = os.path.join(venv_dir, 'bin', 'python') - if {hermetic_venv_scripts!r}: - __re_exec__(venv_python, '-sE', venv_pex) - else: - __re_exec__(venv_python, venv_pex) - - -def __entry_point_from_filename__(filename): - # 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 __execute__: - return entry_point - return os.path.dirname(entry_point) - - -__entry_point__ = None -if '__file__' in locals() and __file__ is not None and os.path.exists(__file__): - __entry_point__ = __entry_point_from_filename__(__file__) -elif '__loader__' in locals(): - 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') - sys.exit(2) - -__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!r}))) - -__venv_dir__ = None -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!r}) - if not ENV.PEX_TOOLS and Variables.PEX_VENV.value_or(ENV, {is_venv!r}): - __venv_dir__ = __maybe_run_venv__( - __entry_point__, - pex_root=__pex_root__, - pex_path=ENV.PEX_PATH or {pex_path!r}, - ) - __entry_point__ = __ensure_pex_installed__( - __entry_point__, pex_root=__pex_root__, pex_hash={pex_hash!r} - ) -else: - os.environ['PEX'] = os.path.realpath(__installed_from__) - -from pex.pex_bootstrapper import bootstrap_pex -bootstrap_pex(__entry_point__, execute=__execute__, venv_dir=__venv_dir__) -""" - - class PEXBuilder(object): """Helper for building PEX environments.""" @@ -595,7 +487,25 @@ def _prepare_code(self): self._pex_info.pex_hash = hashlib.sha1(self._pex_info.dump().encode("utf-8")).hexdigest() self._chroot.write(self._pex_info.dump().encode("utf-8"), PexInfo.PATH, label="manifest") - bootstrap = BOOTSTRAP_ENVIRONMENT.format( + with open(os.path.join(_ABS_PEX_PACKAGE_DIR, "pex_boot.py")) as fp: + pex_boot = fp.read() + + pex_main = dedent( + """ + result = boot( + bootstrap_dir={bootstrap_dir!r}, + pex_root={pex_root!r}, + pex_hash={pex_hash!r}, + has_interpreter_constraints={has_interpreter_constraints!r}, + hermetic_venv_scripts={hermetic_venv_scripts!r}, + pex_path={pex_path!r}, + is_venv={is_venv!r}, + inject_python_args={inject_python_args!r}, + ) + if __SHOULD_EXECUTE__: + sys.exit(result) + """ + ).format( bootstrap_dir=self._pex_info.bootstrap, pex_root=self._pex_info.raw_pex_root, pex_hash=self._pex_info.pex_hash, @@ -603,7 +513,10 @@ def _prepare_code(self): hermetic_venv_scripts=self._pex_info.venv_hermetic_scripts, pex_path=self._pex_info.pex_path, is_venv=self._pex_info.venv, + inject_python_args=self._pex_info.inject_python_args, ) + bootstrap = pex_boot + "\n" + pex_main + self._chroot.write( data=to_bytes(self._shebang + "\n" + self._preamble + "\n" + bootstrap), dst="__main__.py", diff --git a/pex/pex_info.py b/pex/pex_info.py index 6041870d5..6953e50d8 100644 --- a/pex/pex_info.py +++ b/pex/pex_info.py @@ -170,6 +170,16 @@ def build_properties(self, value): self._pex_info["build_properties"] = self.make_build_properties() self._pex_info["build_properties"].update(value) + @property + def inject_python_args(self): + # type: () -> Tuple[str, ...] + return tuple(self._pex_info.get("inject_python_args", ())) + + @inject_python_args.setter + def inject_python_args(self, value): + # type: (Iterable[str]) -> None + self._pex_info["inject_python_args"] = tuple(value) + @property def inject_env(self): # type: () -> Dict[str, str] diff --git a/pex/sh_boot.py b/pex/sh_boot.py index b17d04468..6402a1e75 100644 --- a/pex/sh_boot.py +++ b/pex/sh_boot.py @@ -9,6 +9,7 @@ from textwrap import dedent from pex import dist_metadata, variables +from pex.compatibility import shlex_quote from pex.dist_metadata import Distribution from pex.interpreter import PythonInterpreter, calculate_binary_name from pex.interpreter_constraints import InterpreterConstraints, iter_compatible_versions @@ -20,7 +21,7 @@ from pex.version import __version__ if TYPE_CHECKING: - from typing import Iterable, Optional, Tuple + from typing import Iterable, List, Optional, Tuple import attr # vendor:skip else: @@ -142,13 +143,16 @@ def create_sh_boot_script( `sh` interpreter at `/bin/sh`, and it reduces re-exec overhead to ~2ms in the warm case (and adds ~2ms in the cold case). """ - python = None # type: Optional[str] - python_args = None # type: Optional[str] + python = "" # type: str + python_args = list(pex_info.inject_python_args) # type: List[str] if python_shebang: shebang = python_shebang[2:] if python_shebang.startswith("#!") else python_shebang args = shlex.split(shebang) python = args[0] - python_args = " ".join(args[1:]) + python_args.extend(args[1:]) + venv_python_args = python_args[:] + if pex_info.venv_hermetic_scripts: + venv_python_args.append("-sE") python_names = tuple( _calculate_applicable_binary_names( @@ -181,7 +185,7 @@ def create_sh_boot_script( DEFAULT_PEX_ROOT="$(echo {pex_root})" DEFAULT_PYTHON="{python}" - DEFAULT_PYTHON_ARGS="{python_args}" + PYTHON_ARGS="{python_args}" PEX_ROOT="${{PEX_ROOT:-${{DEFAULT_PEX_ROOT}}}}" INSTALLED_PEX="${{PEX_ROOT}}/{pex_installed_relpath}" @@ -191,12 +195,8 @@ def create_sh_boot_script( # interpreter to use is embedded in the shebang of our venv pex script; so just # execute that script directly. export PEX="$0" - if [ -n "${{VENV_PYTHON_ARGS}}" ]; then - exec "${{INSTALLED_PEX}}/bin/python" "${{VENV_PYTHON_ARGS}}" "${{INSTALLED_PEX}}" \\ - "$@" - else - exec "${{INSTALLED_PEX}}/bin/python" "${{INSTALLED_PEX}}" "$@" - fi + exec "${{INSTALLED_PEX}}/bin/python" ${{VENV_PYTHON_ARGS}} "${{INSTALLED_PEX}}" \\ + "$@" fi find_python() {{ @@ -210,7 +210,7 @@ def create_sh_boot_script( }} if [ -x "${{DEFAULT_PYTHON}}" ]; then - python_exe="${{DEFAULT_PYTHON}} ${{DEFAULT_PYTHON_ARGS}}" + python_exe="${{DEFAULT_PYTHON}}" else python_exe="$(find_python)" fi @@ -223,14 +223,14 @@ def create_sh_boot_script( # __main__.py in our top-level directory; so execute Python against that # directory. export __PEX_EXE__="$0" - exec ${{python_exe}} "${{INSTALLED_PEX}}" "$@" + exec "${{python_exe}}" ${{PYTHON_ARGS}} "${{INSTALLED_PEX}}" "$@" else # The slow path: this PEX zipapp is not installed yet. Run the PEX zipapp so it # can install itself, rebuilding its fast path layout under the PEX_ROOT. if [ -n "${{PEX_VERBOSE:-}}" ]; then echo >&2 "Running zipapp pex to lay itself out under PEX_ROOT." fi - exec ${{python_exe}} "$0" "$@" + exec "${{python_exe}}" ${{PYTHON_ARGS}} "$0" "$@" fi fi @@ -249,9 +249,11 @@ def create_sh_boot_script( ).format( venv="1" if pex_info.venv else "", python=python, - python_args=python_args, + python_args=" ".join(shlex_quote(python_arg) for python_arg in python_args), pythons=" \\\n".join('"{python}"'.format(python=python) for python in python_names), pex_root=pex_info.raw_pex_root, pex_installed_relpath=os.path.relpath(pex_installed_path, pex_info.raw_pex_root), - venv_python_args="-sE" if pex_info.venv_hermetic_scripts else "", + venv_python_args=" ".join( + shlex_quote(venv_python_arg) for venv_python_arg in venv_python_args + ), ) diff --git a/tests/integration/test_inject_python_args.py b/tests/integration/test_inject_python_args.py new file mode 100644 index 000000000..faee435c9 --- /dev/null +++ b/tests/integration/test_inject_python_args.py @@ -0,0 +1,261 @@ +# Copyright 2022 Pex project contributors. +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +import json +import os.path +import shutil +import subprocess +import sys +from textwrap import dedent + +import pytest + +from pex.typing import TYPE_CHECKING +from testing import IS_PYPY, run_pex_command + +if TYPE_CHECKING: + from typing import Any, Iterable, List + +parametrize_execution_mode_args = pytest.mark.parametrize( + "execution_mode_args", + [ + pytest.param([], id="ZIPAPP"), + pytest.param(["--venv"], id="VENV"), + ], +) + + +parametrize_boot_mode_args = pytest.mark.parametrize( + "boot_mode_args", + [ + pytest.param([], id="PYTHON"), + pytest.param(["--sh-boot"], id="SH_BOOT"), + ], +) + + +@pytest.fixture +def exe(tmpdir): + # type: (Any) -> str + exe = os.path.join(str(tmpdir), "exe.py") + with open(exe, "w") as fp: + fp.write( + dedent( + """\ + import json + import sys + import warnings + + + warnings.warn("If you don't eat your meat, you can't have any pudding!") + json.dump(sys.argv[1:], sys.stdout) + """ + ) + ) + return exe + + +def assert_exe_output( + pex, # type: str + warning_expected, # type: bool + prefix_args=(), # type: Iterable[str] +): + process = subprocess.Popen( + args=list(prefix_args) + [pex, "--foo", "bar"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = process.communicate() + error = stderr.decode("utf-8") + assert 0 == process.returncode, error + assert ["--foo", "bar"] == json.loads(stdout) + assert warning_expected == ( + "If you don't eat your meat, you can't have any pudding!" in error + ), error + + +@pytest.mark.skipif( + IS_PYPY and sys.version_info[:2] < (3, 10), + reason=( + "Pex cannot retrieve the original argv when running under PyPy<3.10 which prevents " + "passthrough." + ), +) +@parametrize_execution_mode_args +@parametrize_boot_mode_args +def test_python_args_passthrough( + tmpdir, # type: Any + execution_mode_args, # type: List[str] + boot_mode_args, # type: List[str] + exe, # type: str +): + # type: (...) -> None + + default_shebang_pex = os.path.join(str(tmpdir), "default_shebang.pex") + custom_shebang_pex = os.path.join(str(tmpdir), "custom_shebang.pex") + pex_root = os.path.join(str(tmpdir), "pex_root") + + args = ( + execution_mode_args + + boot_mode_args + + [ + "--pex-root", + pex_root, + "--runtime-pex-root", + pex_root, + "--exe", + exe, + ] + ) + run_pex_command(args=args + ["-o", default_shebang_pex]).assert_success() + run_pex_command( + args=args + + [ + "-o", + custom_shebang_pex, + "--python-shebang", + "{python} -Wignore".format(python=sys.executable), + ] + ).assert_success() + + # N.B.: We execute tests in doubles after a cache nuke to exercise both cold and warm runs + # which take different re-exec paths through the code that all need to preserve Python args. + + # The built-in python shebang args, if any, should be respected. + shutil.rmtree(pex_root) + assert_exe_output(default_shebang_pex, warning_expected=True) + assert_exe_output(default_shebang_pex, warning_expected=True) + assert_exe_output(custom_shebang_pex, warning_expected=False) + assert_exe_output(custom_shebang_pex, warning_expected=False) + + # But they also should be able to be over-ridden. + shutil.rmtree(pex_root) + assert_exe_output( + default_shebang_pex, prefix_args=[sys.executable, "-Wignore"], warning_expected=False + ) + assert_exe_output( + default_shebang_pex, prefix_args=[sys.executable, "-Wignore"], warning_expected=False + ) + assert_exe_output(custom_shebang_pex, prefix_args=[sys.executable], warning_expected=True) + assert_exe_output(custom_shebang_pex, prefix_args=[sys.executable], warning_expected=True) + + +@parametrize_execution_mode_args +@parametrize_boot_mode_args +def test_inject_python_args( + tmpdir, # type: Any + execution_mode_args, # type: List[str] + boot_mode_args, # type: List[str] + exe, # type: str +): + # type: (...) -> None + + pex = os.path.join(str(tmpdir), "pex") + pex_root = os.path.join(str(tmpdir), "pex_root") + + run_pex_command( + args=execution_mode_args + + boot_mode_args + + [ + "--pex-root", + pex_root, + "--runtime-pex-root", + pex_root, + "--exe", + exe, + "--inject-python-args=-W ignore", + "-o", + pex, + ] + ).assert_success() + + assert_exe_output(pex, warning_expected=False) + assert_exe_output(pex, warning_expected=False) + + # N.B.: The original argv cannot be detected by Pex running under PyPy<3.10; so we expect + # warnings to be turned off (the default sealed in by `--inject-python-args`). For all other + # Pythons we support, these explicit command line Python args should be detected and trump the + # injected args by dint of occurring later in the command line. In other words, the command line + # should be as follows and Python is known to pick the last occurrence of the -W option: + # + # python -W ignore -W always ... + # + warning_expected = not IS_PYPY or sys.version_info[:2] >= (3, 10) + assert_exe_output( + pex, prefix_args=[sys.executable, "-W", "always"], warning_expected=warning_expected + ) + assert_exe_output( + pex, prefix_args=[sys.executable, "-W", "always"], warning_expected=warning_expected + ) + + +@pytest.mark.skipif( + sys.version_info[:2] < (3, 10 if IS_PYPY else 9), + reason=( + "The effect of `-u` on the `sys.stdout.buffer` type used in this test is only " + "differentiable from a lack of `-u` for these Pythons." + ), +) +@parametrize_execution_mode_args +@parametrize_boot_mode_args +def test_issue_2422( + tmpdir, # type: Any + execution_mode_args, # type: List[str] + boot_mode_args, # type: List[str] +): + # type: (...) -> None + + pex = os.path.join(str(tmpdir), "pex") + pex_root = os.path.join(str(tmpdir), "pex_root") + + exe = os.path.join(str(tmpdir), "exe.py") + with open(exe, "w") as fp: + fp.write( + dedent( + """\ + import sys + + + if __name__ == "__main__": + print(type(sys.stdout.buffer).__name__) + """ + ) + ) + + args = ["--pex-root", pex_root, "--runtime-pex-root", pex_root, "--exe", exe, "-o", pex] + args.extend(execution_mode_args) + args.extend(boot_mode_args) + + run_pex_command(args=args).assert_success() + shutil.rmtree(pex_root) + assert b"BufferedWriter\n" == subprocess.check_output(args=[sys.executable, exe]) + assert b"BufferedWriter\n" == subprocess.check_output( + args=[pex] + ), "Expected cold run to use buffered io." + assert b"BufferedWriter\n" == subprocess.check_output( + args=[pex] + ), "Expected warm run to use buffered io." + + assert b"FileIO\n" == subprocess.check_output( + args=[sys.executable, "-u", pex] + ), "Expected explicit Python arguments to be honored." + + run_pex_command(args=args + ["--inject-python-args=-u"]).assert_success() + shutil.rmtree(pex_root) + assert b"FileIO\n" == subprocess.check_output(args=[sys.executable, "-u", exe]) + assert b"FileIO\n" == subprocess.check_output( + args=[pex] + ), "Expected cold run to use un-buffered io." + assert b"FileIO\n" == subprocess.check_output( + args=[pex] + ), "Expected warm run to use un-buffered io." + + process = subprocess.Popen( + args=[sys.executable, "-v", pex], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = process.communicate() + assert 0 == process.returncode + assert b"FileIO\n" == stdout, "Expected injected Python arguments to be honored." + assert b"import " in stderr, "Expected explicit Python arguments to be honored as well."