From 18ff5bc217abaa60ad41bae7284b40282cf02630 Mon Sep 17 00:00:00 2001
From: John Sirois <john.sirois@gmail.com>
Date: Tue, 25 Sep 2018 10:51:47 -0600
Subject: [PATCH] Fixup `PEX.demote_bootstrap`: fully unimport.

Previously only root 3rdparty packages copied into the bootstrap were
unimported leaving subpackages in-place. This could lead to a mismatch
in 3rdparty package API expectations when a re-imported 3rdparty root
package from the user PEX sys.path was from a different version of the
3rdparty package.

Fixes #550
---
 pex/pex.py        | 11 +++++++----
 pex/testing.py    |  7 +++++--
 tests/test_pex.py | 27 +++++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/pex/pex.py b/pex/pex.py
index 9cd809af3..2cd0525f1 100644
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -414,11 +414,14 @@ def demote_bootstrap(cls):
     # Move the third party resources pex uses to the end of sys.path for the duration of the run to
     # allow conflicting versions supplied by user dependencies to win during the course of the
     # execution of user code.
-    for _, mod, _ in pkgutil.iter_modules([bootstrap_path]):
-      if mod != root_package:  # We let _pex stay imported
+    for _, submod, _ in pkgutil.iter_modules([bootstrap_path]):
+      if submod != root_package:  # We let _pex stay imported
         TRACER.log('Un-importing third party bootstrap dependency %s from %s'
-                   % (mod, bootstrap_path))
-        sys.modules.pop(mod)
+                   % (submod, bootstrap_path))
+        sys.modules.pop(submod)
+        submod_prefix = submod + '.'
+        for submod in [m for m in sys.modules.keys() if m.startswith(submod_prefix)]:
+          sys.modules.pop(submod)
     sys.path.pop(bootstrap_path_index)
     sys.path.append(bootstrap_path)
 
diff --git a/pex/testing.py b/pex/testing.py
index e016cb82a..2146c9eb1 100644
--- a/pex/testing.py
+++ b/pex/testing.py
@@ -180,7 +180,7 @@ def make_bdist(name='my_project', version='0.0.0', installer_impl=EggInstaller,
 """
 
 
-def write_simple_pex(td, exe_contents, dists=None, sources=None, coverage=False):
+def write_simple_pex(td, exe_contents, dists=None, sources=None, coverage=False, interpreter=None):
   """Write a pex file that contains an executable entry point
 
   :param td: temporary directory path
@@ -189,6 +189,7 @@ def write_simple_pex(td, exe_contents, dists=None, sources=None, coverage=False)
   :param dists: distributions to include, typically sdists or bdists
   :param sources: sources to include, as a list of pairs (env_filename, contents)
   :param coverage: include coverage header
+  :param interpreter: a custom interpreter to use to build the pex
   """
   dists = dists or []
   sources = sources or []
@@ -198,7 +199,9 @@ def write_simple_pex(td, exe_contents, dists=None, sources=None, coverage=False)
   with open(os.path.join(td, 'exe.py'), 'w') as fp:
     fp.write(exe_contents)
 
-  pb = PEXBuilder(path=td, preamble=COVERAGE_PREAMBLE if coverage else None)
+  pb = PEXBuilder(path=td,
+                  preamble=COVERAGE_PREAMBLE if coverage else None,
+                  interpreter=interpreter)
 
   for dist in dists:
     pb.add_dist_location(dist.location)
diff --git a/tests/test_pex.py b/tests/test_pex.py
index 156439a7f..6b22b70f5 100644
--- a/tests/test_pex.py
+++ b/tests/test_pex.py
@@ -6,6 +6,7 @@
 import sys
 import textwrap
 from contextlib import contextmanager
+from pex.interpreter import PythonInterpreter
 from types import ModuleType
 
 import pytest
@@ -373,3 +374,29 @@ def test_pex_run_custom_setuptools_useable():
       )
       rc = PEX(pex.path()).run()
       assert rc == 0
+
+
+def test_pex_run_conflicting_custom_setuptools_useable():
+  # Here we use an older setuptools to build the pex which has a newer setuptools requirement.
+  # These setuptools dists have different pkg_resources APIs:
+  # $ diff \
+  #   <(zipinfo -1 setuptools-20.3.1-py2.py3-none-any.whl | grep pkg_resources/ | sort) \
+  #   <(zipinfo -1 setuptools-40.4.3-py2.py3-none-any.whl | grep pkg_resources/ | sort)
+  # 2a3,4
+  # > pkg_resources/py31compat.py
+  # > pkg_resources/_vendor/appdirs.py
+  with temporary_dir() as resolve_cache:
+    dists = resolve(['setuptools==20.3.1'], cache=resolve_cache)
+    interpreter = PythonInterpreter.from_binary(sys.executable,
+                                                path_extras=[dist.location for dist in dists],
+                                                include_site_extras=False)
+    dists = resolve(['setuptools==40.4.3'], cache=resolve_cache)
+    with temporary_dir() as temp_dir:
+      pex = write_simple_pex(
+        temp_dir,
+        'from pkg_resources import appdirs, py31compat',
+        dists=dists,
+        interpreter=interpreter
+      )
+      rc = PEX(pex.path()).run()
+      assert rc == 0