diff --git a/pex/bin/pex.py b/pex/bin/pex.py index 7343e4241..406b74e14 100755 --- a/pex/bin/pex.py +++ b/pex/bin/pex.py @@ -231,6 +231,17 @@ def configure_clp_pex_options(parser): "problems with tools or libraries that are confused by symlinked source files." ), ) + group.add_argument( + "--non-hermetic-venv-scripts", + dest="venv_hermetic_scripts", + action="store_false", + default=True, + help=( + "If --venv is specified, don't rewrite Python script shebangs in the venv to pass " + "`-sE` to the interpreter; for example, to enable running the venv PEX itself or its " + "Python scripts with a custom `PYTHONPATH`." + ), + ) group.add_argument( "--always-write-cache", @@ -640,6 +651,7 @@ def build_pex( pex_info.venv_bin_path = options.venv or BinPath.FALSE pex_info.venv_copies = options.venv_copies pex_info.venv_site_packages_copies = options.venv_site_packages_copies + pex_info.venv_hermetic_scripts = options.venv_hermetic_scripts pex_info.includes_tools = options.include_tools or options.venv pex_info.pex_path = options.pex_path.split(os.pathsep) if options.pex_path else () pex_info.ignore_errors = options.ignore_errors diff --git a/pex/pex_bootstrapper.py b/pex/pex_bootstrapper.py index 3e168cbde..31d19e7e4 100644 --- a/pex/pex_bootstrapper.py +++ b/pex/pex_bootstrapper.py @@ -447,6 +447,7 @@ def _bootstrap(entry_point): @attr.s(frozen=True) class VenvPex(object): venv_dir = attr.ib() # type: str + hermetic_scripts = attr.ib(default=True) # type: bool pex = attr.ib(init=False) # type: str python = attr.ib(init=False) # type: str @@ -461,7 +462,12 @@ def __attrs_post_init__(self): def execute_args(self, *additional_args): # type: (*str) -> List[str] - return [self.python, "-sE", self.pex] + list(additional_args) + argv = [self.python] + 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 @@ -545,6 +551,7 @@ def ensure_venv( ), collisions_ok=collisions_ok, symlink=symlink, + hermetic_scripts=pex_info.venv_hermetic_scripts, ) # There are popular Linux distributions with shebang length limits @@ -567,7 +574,7 @@ def ensure_venv( break - return VenvPex(venv_dir) + return VenvPex(venv_dir, hermetic_scripts=pex_info.venv_hermetic_scripts) # NB: This helper is used by the PEX bootstrap __main__.py as well as the __pex__/__init__.py diff --git a/pex/pex_builder.py b/pex/pex_builder.py index 036758aa3..0663950d9 100644 --- a/pex/pex_builder.py +++ b/pex/pex_builder.py @@ -100,7 +100,10 @@ def __maybe_run_venv__(pex, pex_root, pex_path): TRACER.log('Executing venv PEX for {{}} at {{}}'.format(pex, venv_pex)) venv_python = os.path.join(venv_dir, 'bin', 'python') - __re_exec__(venv_python, '-sE', venv_pex) + 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): @@ -521,6 +524,7 @@ def _prepare_code(self): pex_root=self._pex_info.raw_pex_root, pex_hash=self._pex_info.pex_hash, has_interpreter_constraints=bool(self._pex_info.interpreter_constraints), + hermetic_venv_scripts=self._pex_info.venv_hermetic_scripts, pex_path=self._pex_info.pex_path, is_venv=self._pex_info.venv, ) diff --git a/pex/pex_info.py b/pex/pex_info.py index dfb462b44..e99700d4f 100644 --- a/pex/pex_info.py +++ b/pex/pex_info.py @@ -231,6 +231,16 @@ def venv_site_packages_copies(self, value): # type: (bool) -> None self._pex_info["venv_site_packages_copies"] = value + @property + def venv_hermetic_scripts(self): + # type: () -> bool + return self._pex_info.get("venv_hermetic_scripts", True) + + @venv_hermetic_scripts.setter + def venv_hermetic_scripts(self, value): + # type: (bool) -> None + self._pex_info["venv_hermetic_scripts"] = value + def _venv_dir( self, pex_root, # type: str diff --git a/pex/sh_boot.py b/pex/sh_boot.py index 556e78d8c..1e0b4e49f 100644 --- a/pex/sh_boot.py +++ b/pex/sh_boot.py @@ -175,6 +175,7 @@ def create_sh_boot_script( set -eu VENV="{venv}" + VENV_PYTHON_ARGS="{venv_python_args}" # N.B.: This ensures tilde-expansion of the DEFAULT_PEX_ROOT value. DEFAULT_PEX_ROOT="$(echo {pex_root})" @@ -190,7 +191,12 @@ 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" - exec "${{INSTALLED_PEX}}/bin/python" -sE "${{INSTALLED_PEX}}" "$@" + 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 fi find_python() {{ @@ -247,4 +253,5 @@ def create_sh_boot_script( 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 "", ) diff --git a/pex/tools/commands/venv.py b/pex/tools/commands/venv.py index 608d94a3b..12c4b5054 100644 --- a/pex/tools/commands/venv.py +++ b/pex/tools/commands/venv.py @@ -172,8 +172,9 @@ def add_arguments(cls, parser): action="store_false", default=True, help=( - "Don't rewrite console script shebangs in the venv to pass `-sE` to the interpreter; " - "for example, to enable running venv scripts with a custom `PYTHONPATH`." + "Don't rewrite Python script shebangs in the venv to pass `-sE` to the " + "interpreter; for example, to enable running the venv PEX itself or its Python " + "scripts with a custom `PYTHONPATH`." ), ) cls.register_global_arguments(parser, include_verbosity=False) diff --git a/pex/venv/pex.py b/pex/venv/pex.py index b0b484ed3..4bbe9e3c7 100644 --- a/pex/venv/pex.py +++ b/pex/venv/pex.py @@ -95,6 +95,11 @@ class CollisionError(Exception): """Indicates multiple distributions provided the same file when merging a PEX into a venv.""" +def _script_python_args(hermetic): + # type: (bool) -> Optional[str] + return "-sE" if hermetic else None + + def populate_venv( venv, # type: Virtualenv pex, # type: PEX @@ -108,7 +113,12 @@ def populate_venv( # type: (...) -> str venv_python = python or venv.interpreter.binary - shebang = "#!{} -sE".format(venv_python) + + shebang_argv = [venv_python] + python_args = _script_python_args(hermetic=hermetic_scripts) + if python_args: + shebang_argv.append(python_args) + shebang = "#!{shebang}".format(shebang=" ".join(shebang_argv)) provenance = defaultdict(list) @@ -276,8 +286,9 @@ def _populate_deps( print(rel_extra_path, file=fp) # 3. Re-write any (console) scripts to use the venv Python. - script_python_args = "-sE" if hermetic_scripts else None - for script in venv.rewrite_scripts(python=venv_python, python_args=script_python_args): + for script in venv.rewrite_scripts( + python=venv_python, python_args=_script_python_args(hermetic=hermetic_scripts) + ): TRACER.log("Re-writing {}".format(script)) @@ -375,7 +386,11 @@ def sys_executable_paths(): ): sys.stderr.write("Re-execing from {{}}\\n".format(sys.executable)) os.environ[current_interpreter_blessed_env_var] = "1" - os.execv(python, [python, "-sE"] + sys.argv) + argv = [python] + if {hermetic_re_exec!r}: + argv.append("-sE") + argv.extend(sys.argv) + os.execv(python, argv) pex_file = os.environ.get("PEX", None) if pex_file: @@ -624,6 +639,7 @@ def sys_executable_paths(): if venv.interpreter.version[0] == 2 else "exec(ast, globals_map, locals_map)" ), + hermetic_re_exec=pex_info.venv_hermetic_scripts, ) ) with open(venv.join_path("__main__.py"), "w") as fp: diff --git a/tests/integration/venv_ITs/test_issue_2065.py b/tests/integration/venv_ITs/test_issue_2065.py new file mode 100644 index 000000000..89d810075 --- /dev/null +++ b/tests/integration/venv_ITs/test_issue_2065.py @@ -0,0 +1,154 @@ +# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +import json +import os +import subprocess +from textwrap import dedent + +import pytest + +from pex.common import safe_open +from pex.testing import make_env, run_pex_command +from pex.typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, List + + +@pytest.mark.parametrize( + ["boot_args"], + [ + pytest.param([], id="__main__.py boot"), + pytest.param(["--sh-boot"], id="--sh-boot"), + ], +) +def test_venv_pex_script_non_hermetic( + tmpdir, # type: Any + boot_args, # type: List[str] +): + # type: (...) -> None + + # A console script that injects an element in the PYTHONPATH. + ot_simulator_src = os.path.join(str(tmpdir), "src") + with safe_open(os.path.join(ot_simulator_src, "ot_simulator.py"), "w") as fp: + fp.write( + dedent( + """\ + import os + import sys + + def run(): + pythonpath = ["injected"] + existing_pythonpath = os.environ.get("PYTHONPATH") + if existing_pythonpath: + pythonpath.extend(existing_pythonpath.split(os.pathsep)) + os.environ["PYTHONPATH"] = os.pathsep.join(pythonpath) + + os.execv(sys.argv[1], sys.argv[1:]) + """ + ) + ) + with safe_open(os.path.join(ot_simulator_src, "setup.cfg"), "w") as fp: + fp.write( + dedent( + """\ + [metadata] + name = ot-simulator + version = 0.0.1 + + [options] + py_modules = + ot_simulator + + [options.entry_points] + console_scripts = + instrument = ot_simulator:run + """ + ) + ) + with safe_open(os.path.join(ot_simulator_src, "setup.py"), "w") as fp: + fp.write("from setuptools import setup; setup()") + + # An entrypoint that can observe the PYTHONPATH / sys.path. + app = os.path.join(str(tmpdir), "app.exe") + with safe_open(app, "w") as fp: + fp.write( + dedent( + """\ + import json + import os + import sys + + json.dump( + { + "PYTHONPATH": os.environ.get("PYTHONPATH"), + "sys.path": sys.path + }, + sys.stdout + ) + """ + ) + ) + + pex_root = os.path.join(str(tmpdir), "pex_root") + + def create_app_pex(hermetic_scripts): + # type: (bool) -> str + pex = os.path.join( + str(tmpdir), "{}-app.pex".format("hermetic" if hermetic_scripts else "non-hermetic") + ) + argv = [ + "--pex-root", + pex_root, + "--runtime-pex-root", + pex_root, + ot_simulator_src, + "--exe", + app, + "--venv", + "-o", + pex, + ] + boot_args + if not hermetic_scripts: + argv.append("--non-hermetic-venv-scripts") + run_pex_command(argv).assert_success() + return pex + + cwd = os.path.join(str(tmpdir), "cwd") + os.mkdir(cwd) + + # A standard hermetic venv pex should be able to see PYTHONPATH but not have its sys.path + # tainted by it. + hermetic_app_pex = create_app_pex(hermetic_scripts=True) + hermetic = json.loads( + subprocess.check_output( + args=[hermetic_app_pex], cwd=cwd, env=make_env(PYTHONPATH="ambient") + ) + ) + assert "ambient" == hermetic["PYTHONPATH"] + assert os.path.join(cwd, "ambient") not in hermetic["sys.path"] + + # A non-hermetic venv pex should be able to both see PYTHONPATH and have it affect its sys.path. + non_hermetic_app_pex = create_app_pex(hermetic_scripts=False) + baseline = json.loads( + subprocess.check_output( + args=[non_hermetic_app_pex], cwd=cwd, env=make_env(PYTHONPATH="ambient") + ) + ) + assert "ambient" == baseline["PYTHONPATH"] + assert os.path.join(cwd, "ambient") in baseline["sys.path"] + + # A non-hermetic venv pex should have the non-hermeticity extend to its console scripts in + # addition to the main entry point `pex` script. + instrumented = json.loads( + subprocess.check_output( + args=[non_hermetic_app_pex, non_hermetic_app_pex], + cwd=cwd, + env=make_env(PYTHONPATH="ambient", PEX_SCRIPT="instrument"), + ) + ) + assert "injected:ambient" == instrumented["PYTHONPATH"] + assert sorted(baseline["sys.path"] + [os.path.join(cwd, "injected")]) == sorted( + instrumented["sys.path"] + )