From 4942bcd303cbd0698dde885a84c6464bfd919edf Mon Sep 17 00:00:00 2001 From: John Sirois Date: Mon, 24 Jul 2023 07:47:12 -0600 Subject: [PATCH] Cleanup `sys.path` after `__pex__` is imported. (#2189) This fixes #1954 by ensuring all vendored code is uninstalled from the path and bootstrap code is demoted to the end of `sys.path` after `__pex__` is imported. Fixes #1954 --------- Co-authored-by: Zameer Manji --- pex/bootstrap.py | 18 ++++-- pex/pex_bootstrapper.py | 9 +++ tests/integration/test_pex_import.py | 82 +++++++++++++++++++++++----- 3 files changed, 89 insertions(+), 20 deletions(-) diff --git a/pex/bootstrap.py b/pex/bootstrap.py index 1634b0340..03cb57fe5 100644 --- a/pex/bootstrap.py +++ b/pex/bootstrap.py @@ -6,6 +6,7 @@ import os import sys +import types class Bootstrap(object): @@ -43,7 +44,7 @@ def path(self): # type: () -> str return self._sys_path_entry - def demote(self): + def demote(self, disable_vendor_importer=True): """Demote the bootstrap code to the end of the `sys.path` so it is found last. :return: The list of un-imported bootstrap modules. @@ -60,9 +61,10 @@ def demote(self): unimported_modules = [] for name, module in reversed(sorted(sys.modules.items())): + if "pex.third_party" == name and not disable_vendor_importer: + continue if self.imported_from_bootstrap(module): unimported_modules.append(sys.modules.pop(name)) - return unimported_modules def imported_from_bootstrap(self, module): @@ -73,6 +75,12 @@ def imported_from_bootstrap(self, module): :rtype: bool """ + # Python 2.7 does some funky imports in the email stdlib package that cause havoc with + # un-importing. Since all our own importing just goes through the vanilla importers we can + # safely ignore all but the standard module type. + if not isinstance(module, types.ModuleType): + return False + # A vendored module. path = getattr(module, "__file__", None) if path and os.path.realpath(path).startswith(self._realpath): @@ -93,8 +101,8 @@ def __repr__(self): ) -def demote(): - # type: () -> None +def demote(disable_vendor_importer=True): + # type: (bool) -> None """Demote PEX bootstrap code to the end of `sys.path` and uninstall all PEX vendored code.""" from . import third_party @@ -115,7 +123,7 @@ def log(msg, V=1): bootstrap = Bootstrap.locate() log("Demoting code from %s" % bootstrap, V=2) - for module in bootstrap.demote(): + for module in bootstrap.demote(disable_vendor_importer=disable_vendor_importer): log("un-imported {}".format(module), V=9) import pex diff --git a/pex/pex_bootstrapper.py b/pex/pex_bootstrapper.py index abe353dbb..ce0583cc2 100644 --- a/pex/pex_bootstrapper.py +++ b/pex/pex_bootstrapper.py @@ -601,6 +601,15 @@ def bootstrap_pex( VendorImporter.install( uninstallable=False, prefix="__pex__", path_items=["."], root=location ) + + from pex import bootstrap + + # For inscrutable reasons, CPython 2.7 (not PyPy 2.7 and not any other CPython or PyPy + # version supported by Pex) cannot handle pex.third_party being un-imported at this + # stage; so we just skip that cleanup step for every interpreter to keep things simple. + # The main point here is that we've un-imported all "exposed" Pex vendored code, notably + # `attrs`. + bootstrap.demote(disable_vendor_importer=False) return interpreter_test = InterpreterTest(entry_point=entry_point, pex_info=pex_info) diff --git a/tests/integration/test_pex_import.py b/tests/integration/test_pex_import.py index 79e15e270..5bb15f161 100644 --- a/tests/integration/test_pex_import.py +++ b/tests/integration/test_pex_import.py @@ -3,20 +3,20 @@ import os.path import subprocess -import sys from textwrap import dedent import colors import pytest +from pex import targets from pex.common import safe_open -from pex.interpreter import PythonInterpreter from pex.layout import DEPS_DIR, Layout from pex.resolve.pex_repository_resolver import resolve_from_pex from pex.targets import Targets from pex.testing import make_env, run_pex_command from pex.typing import TYPE_CHECKING from pex.variables import ENV +from pex.venv.virtualenv import Virtualenv if TYPE_CHECKING: from typing import Any, List, Text @@ -35,6 +35,12 @@ def test_import_from_pex( ): # type: (...) -> None + empty_env_dir = os.path.join(str(tmpdir), "empty_env") + empty_venv = Virtualenv.create( + venv_dir=empty_env_dir, + ) + empty_python = empty_venv.interpreter.binary + src = os.path.join(str(tmpdir), "src") with safe_open(os.path.join(src, "first_party.py"), "w") as fp: fp.write( @@ -62,6 +68,8 @@ def warn(msg): "-D", src, "ansicolors==1.1.8", + # Add pex to verify that it will shadow bootstrap pex + "pex==2.1.139", "-o", pex, "--layout", @@ -73,18 +81,28 @@ def warn(msg): def execute_with_pex_on_pythonpath(code): # type: (str) -> Text return ( - subprocess.check_output(args=[sys.executable, "-c", code], env=make_env(PYTHONPATH=pex)) + subprocess.check_output( + args=[empty_python, "-c", code], env=make_env(PYTHONPATH=pex), cwd=str(tmpdir) + ) .decode("utf-8") .strip() ) + def get_third_party_prefix(): + if is_venv: + return os.path.join(pex_root, "venvs") + elif layout is Layout.LOOSE: + return os.path.join(pex, DEPS_DIR) + else: + return os.path.join(pex_root, "installed_wheels") + # Verify 3rd party code can be imported hermetically from the PEX. alternate_pex_root = os.path.join(str(tmpdir), "alternate_pex_root") with ENV.patch(PEX_ROOT=alternate_pex_root): ambient_sys_path = [ installed_distribution.fingerprinted_distribution.distribution.location for installed_distribution in resolve_from_pex( - targets=Targets(interpreters=(PythonInterpreter.from_binary(sys.executable),)), + targets=Targets.from_target(targets.current()), pex=pex, requirements=["ansicolors==1.1.8"], ).installed_distributions @@ -95,24 +113,19 @@ def execute_with_pex_on_pythonpath(code): """\ # Executor code like the AWS runtime. import sys - + sys.path = {ambient_sys_path!r} + sys.path - + # User code residing in the PEX. from __pex__ import colors - + print(colors.__file__) """.format( ambient_sys_path=ambient_sys_path ) ) ) - if is_venv: - expected_prefix = os.path.join(pex_root, "venvs") - elif layout is Layout.LOOSE: - expected_prefix = os.path.join(pex, DEPS_DIR) - else: - expected_prefix = os.path.join(pex_root, "installed_wheels") + expected_prefix = get_third_party_prefix() assert third_party_path.startswith( expected_prefix ), "Expected 3rd party ansicolors path {path} to start with {expected_prefix}".format( @@ -141,10 +154,49 @@ def execute_with_pex_on_pythonpath(code): import colors import first_party - - + + print(colors.blue("42")) first_party.warn("Vogon") """ ) ) + + # Verify bootstrap code does not leak attrs from vendored code + assert "no leak" == execute_with_pex_on_pythonpath( + dedent( + """\ + import __pex__ + + try: + import attr + print(attr.__file__) + except ImportError: + print("no leak") + """ + ) + ) + + # Verify bootstrap pex code is demoted to the end of `sys.path` + assert os.path.join(pex, ".bootstrap") == execute_with_pex_on_pythonpath( + dedent( + """\ + import __pex__ + import sys + print(sys.path[-1]) + """ + ) + ) + + # Verify third party pex shadows bootstrap pex + pex_third_party_path = execute_with_pex_on_pythonpath( + dedent( + """\ + import __pex__ + import pex + print(pex.__file__) + """ + ) + ) + + assert pex_third_party_path.startswith(get_third_party_prefix())