Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Feb 26, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions pex/bin/pex.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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
Expand Down
11 changes: 9 additions & 2 deletions pex/pex_bootstrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
6 changes: 5 additions & 1 deletion pex/pex_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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,
)
Expand Down
10 changes: 10 additions & 0 deletions pex/pex_info.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 8 additions & 1 deletion pex/sh_boot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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})"
Expand All @@ -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}}" "$@"
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had missed this 1 spot. That's the last.

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() {{
Expand Down Expand Up @@ -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 "",
)
5 changes: 3 additions & 2 deletions pex/tools/commands/venv.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
24 changes: 20 additions & 4 deletions pex/venv/pex.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)

Expand Down Expand Up @@ -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))


Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
154 changes: 154 additions & 0 deletions tests/integration/venv_ITs/test_issue_2065.py
Original file line number Diff line number Diff line change
@@ -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")
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As named, this simulates OpenTelemetry auto-instrumentation - actually turning the manual test detailed here #2065 (comment) into an IT seemed too fragile / likely flaky.

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