Introduce --non-hermetic-venv-scripts. (#2068)

This provides the same option when building a PEX as the
`--non-hermetic-scripts` does when creating a venv from a PEX using
`pex-tools`. In addition, both means of creating a non-hermetic venv
are extended to the venv `pex` script itself.

Closes #2065

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

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

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,
         )

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

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 "",
     )

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)

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:

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"]
+    )