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 <[email protected]>

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

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)

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