Repository: pex-tool/pex

Issue:
pex shadows pkg_resources, which can break newer versions of setuptools
with a version of setuptools inside pex's allowed range:

```
[omerta pants (kwlzn/pantsd_bugfixes)]$ pex "setuptools<34"
Python 2.7.10 (default, Dec 16 2015, 14:09:45) 
[GCC 4.2.1 Compatible Apple LLVM 7.0.2 (clang-700.1.81)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
(InteractiveConsole)
>>> from setuptools.sandbox import run_setup
>>> 
```

with a newer one:

```
[omerta pants (kwlzn/pantsd_bugfixes)]$ pex setuptools==36.2.7
Python 2.7.10 (default, Dec 16 2015, 14:09:45) 
[GCC 4.2.1 Compatible Apple LLVM 7.0.2 (clang-700.1.81)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
(InteractiveConsole)
>>> from setuptools.sandbox import run_setup
Traceback (most recent call last):
  File "<console>", line 1, in <module>
  File "/private/var/folders/bb/msvj_4_s741g93nzgrj19twc0000gn/T/tmp1zkDfP/.deps/setuptools-36.2.7-py2.py3-none-any.whl/setuptools/sandbox.py", line 15, in <module>
    import pkg_resources.py31compat
ImportError: No module named py31compat
>>> 
```

because:

```
>>> import pkg_resources
>>> pkg_resources.__file__
'/private/var/folders/bb/msvj_4_s741g93nzgrj19twc0000gn/T/tmpsTSy8w/.bootstrap/pkg_resources/__init__.pyc'
>>> 
```

poc:

```
[omerta pants (kwlzn/pantsd_bugfixes)]$ pex setuptools==36.2.7
Python 2.7.10 (default, Dec 16 2015, 14:09:45) 
[GCC 4.2.1 Compatible Apple LLVM 7.0.2 (clang-700.1.81)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
(InteractiveConsole)
>>> import sys
>>> sys.path
['/private/var/folders/bb/msvj_4_s741g93nzgrj19twc0000gn/T/tmpCoO7yi/.bootstrap', '/var/folders/bb/msvj_4_s741g93nzgrj19twc0000gn/T/tmpCoO7yi', '/opt/ee/python/2.7/lib/python27.zip', '/opt/ee/python/2.7/lib/python2.7', '/opt/ee/python/2.7/lib/python2.7/plat-darwin', '/opt/ee/python/2.7/lib/python2.7/plat-mac', '/opt/ee/python/2.7/lib/python2.7/plat-mac/lib-scriptpackages', '/opt/ee/python/2.7/lib/python2.7/lib-tk', '/opt/ee/python/2.7/lib/python2.7/lib-old', '/opt/ee/python/2.7/lib/python2.7/lib-dynload', '/private/var/folders/bb/msvj_4_s741g93nzgrj19twc0000gn/T/tmpCoO7yi/.deps/setuptools-36.2.7-py2.py3-none-any.whl']
>>> sys.path.insert(0, sys.path[-1])
>>> del sys.modules['pkg_resources']
>>> from setuptools.sandbox import run_setup
>>> 
```


Discussion:
We probably need to vendor our `.bootstrap/` code up to and including pex itself on the off chance a user wants to create a pex of their own code that uses pex as a library.
Or... delegate to the final picked runner in a contextmanager that removes .bootstrap from the sys.path and the re-adds it on exit so pex exit handlers - if any can run - which is the POC above. I'll try this first.

Resolving patch:
diff --git a/pex/pex.py b/pex/pex.py
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -4,8 +4,8 @@
 from __future__ import absolute_import, print_function
 
 import os
+import pkgutil
 import sys
-from contextlib import contextmanager
 from distutils import sysconfig
 from site import USER_SITE
 
@@ -249,28 +249,16 @@ def minimum_sys(cls, inherit_path):
     return sys_path, sys_path_importer_cache, sys_modules
 
   @classmethod
-  @contextmanager
   def patch_pkg_resources(cls, working_set):
     """Patch pkg_resources given a new working set."""
-    def patch(working_set):
-      pkg_resources.working_set = working_set
-      pkg_resources.require = working_set.require
-      pkg_resources.iter_entry_points = working_set.iter_entry_points
-      pkg_resources.run_script = pkg_resources.run_main = working_set.run_script
-      pkg_resources.add_activation_listener = working_set.subscribe
-
-    old_working_set = pkg_resources.working_set
-    patch(working_set)
-    try:
-      yield
-    finally:
-      patch(old_working_set)
-
-  # Thar be dragons -- when this contextmanager exits, the interpreter is
-  # potentially in a wonky state since the patches here (minimum_sys_modules
-  # for example) actually mutate global state.  This should not be
-  # considered a reversible operation despite being a contextmanager.
-  @contextmanager
+    pkg_resources.working_set = working_set
+    pkg_resources.require = working_set.require
+    pkg_resources.iter_entry_points = working_set.iter_entry_points
+    pkg_resources.run_script = pkg_resources.run_main = working_set.run_script
+    pkg_resources.add_activation_listener = working_set.subscribe
+
+  # Thar be dragons -- when this function exits, the interpreter is potentially in a wonky state
+  # since the patches here (minimum_sys_modules for example) actually mutate global state.
   def patch_sys(self, inherit_path):
     """Patch sys with all site scrubbed."""
     def patch_dict(old_value, new_value):
@@ -282,14 +270,11 @@ def patch_all(path, path_importer_cache, modules):
       patch_dict(sys.path_importer_cache, path_importer_cache)
       patch_dict(sys.modules, modules)
 
-    old_sys_path, old_sys_path_importer_cache, old_sys_modules = (
-        sys.path[:], sys.path_importer_cache.copy(), sys.modules.copy())
     new_sys_path, new_sys_path_importer_cache, new_sys_modules = self.minimum_sys(inherit_path)
 
     new_sys_path.extend(merge_split(self._pex_info.pex_path, self._vars.PEX_PATH))
 
     patch_all(new_sys_path, new_sys_path_importer_cache, new_sys_modules)
-    yield
 
   def _wrap_coverage(self, runner, *args):
     if not self._vars.PEX_COVERAGE and self._vars.PEX_COVERAGE_FILENAME is None:
@@ -360,14 +345,10 @@ def execute(self):
       pex_inherit_path = self._vars.PEX_INHERIT_PATH
       if pex_inherit_path == "false":
         pex_inherit_path = self._pex_info.inherit_path
-      with self.patch_sys(pex_inherit_path):
-        working_set = self._activate()
-        TRACER.log('PYTHONPATH contains:')
-        for element in sys.path:
-          TRACER.log('  %c %s' % (' ' if os.path.exists(element) else '*', element))
-        TRACER.log('  * - paths that do not exist or will be imported via zipimport')
-        with self.patch_pkg_resources(working_set):
-          self._wrap_coverage(self._wrap_profiling, self._execute)
+      self.patch_sys(pex_inherit_path)
+      working_set = self._activate()
+      self.patch_pkg_resources(working_set)
+      self._wrap_coverage(self._wrap_profiling, self._execute)
     except Exception:
       # Allow the current sys.excepthook to handle this app exception before we tear things down in
       # finally, then reraise so that the exit status is reflected correctly.
@@ -415,6 +396,37 @@ def _execute(self):
       TRACER.log('No entry point specified, dropping into interpreter')
       return self.execute_interpreter()
 
+  @classmethod
+  def demote_bootstrap(cls):
+    TRACER.log('Bootstrap complete, performing final sys.path modifications...')
+
+    bootstrap_path = __file__
+    module_import_path = __name__.split('.')
+    root_package = module_import_path[0]
+
+    # For example, our __file__ might be requests.pex/.bootstrap/_pex/pex.pyc and our import path
+    # _pex.pex; so we walk back through all the module components of our import path to find the
+    # base sys.path entry where we were found (requests.pex/.bootstrap in this example).
+    for _ in module_import_path:
+      bootstrap_path = os.path.dirname(bootstrap_path)
+    bootstrap_path_index = sys.path.index(bootstrap_path)
+
+    # 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
+        TRACER.log('Un-importing third party bootstrap dependency %s from %s'
+                   % (mod, bootstrap_path))
+        sys.modules.pop(mod)
+    sys.path.pop(bootstrap_path_index)
+    sys.path.append(bootstrap_path)
+
+    TRACER.log('PYTHONPATH contains:')
+    for element in sys.path:
+      TRACER.log('  %c %s' % (' ' if os.path.exists(element) else '*', element))
+    TRACER.log('  * - paths that do not exist or will be imported via zipimport')
+
   def execute_interpreter(self):
     if sys.argv[1:]:
       program = sys.argv[1]
@@ -429,6 +441,8 @@ def execute_interpreter(self):
       sys.argv = sys.argv[1:]
       self.execute_content(program, content)
     else:
+      self.demote_bootstrap()
+
       import code
       code.interact()
 
@@ -447,6 +461,8 @@ def execute_script(self, script_name):
 
   @classmethod
   def execute_content(cls, name, content, argv0=None):
+    cls.demote_bootstrap()
+
     argv0 = argv0 or name
     try:
       ast = compile(content, name, 'exec', flags=0, dont_inherit=1)
@@ -474,13 +490,17 @@ def execute_entry(cls, entry_point):
     runner = cls.execute_pkg_resources if ':' in entry_point else cls.execute_module
     return runner(entry_point)
 
-  @staticmethod
-  def execute_module(module_name):
+  @classmethod
+  def execute_module(cls, module_name):
+    cls.demote_bootstrap()
+
     import runpy
     runpy.run_module(module_name, run_name='__main__')
 
-  @staticmethod
-  def execute_pkg_resources(spec):
+  @classmethod
+  def execute_pkg_resources(cls, spec):
+    cls.demote_bootstrap()
+
     entry = EntryPoint.parse("run = {0}".format(spec))
 
     # See https://pythonhosted.org/setuptools/history.html#id25 for rationale here.


Test patch:
diff --git a/pex/testing.py b/pex/testing.py
--- a/pex/testing.py
+++ b/pex/testing.py
@@ -201,7 +201,7 @@ def write_simple_pex(td, exe_contents, dists=None, sources=None, coverage=False)
   pb = PEXBuilder(path=td, preamble=COVERAGE_PREAMBLE if coverage else None)
 
   for dist in dists:
-    pb.add_egg(dist.location)
+    pb.add_dist_location(dist.location)
 
   for env_filename, contents in sources:
     src_path = os.path.join(td, env_filename)
diff --git a/tests/test_integration.py b/tests/test_integration.py
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -1001,3 +1001,76 @@ def test_multiplatform_entrypoint():
 
     greeting = subprocess.check_output([pex_out_path])
     assert b'Hello World!' == greeting.strip()
+
+
+@contextmanager
+def pex_with_entrypoints(entry_point):
+  setup_py = dedent("""
+    from setuptools import setup
+
+    setup(
+      name='my_app',
+      version='0.0.0',
+      zip_safe=True,
+      packages=[''],
+      install_requires=['setuptools==36.2.7'],
+      entry_points={'console_scripts': ['my_app_function = my_app:do_something',
+                                        'my_app_module = my_app']},
+    )
+  """)
+
+  my_app = dedent("""
+    from setuptools.sandbox import run_setup
+
+    def do_something():
+      return run_setup
+
+    if __name__ == '__main__':
+      do_something()
+  """)
+
+  with temporary_content({'setup.py': setup_py, 'my_app.py': my_app}) as project_dir:
+    with temporary_dir() as out:
+      pex = os.path.join(out, 'pex.pex')
+      pex_command = ['--validate-entry-point', '-c', entry_point, project_dir, '-o', pex]
+      results = run_pex_command(pex_command)
+      results.assert_success()
+      yield pex
+
+
+def test_pex_script_module_custom_setuptools_useable():
+  with pex_with_entrypoints('my_app_module') as pex:
+    stdout, rc = run_simple_pex(pex, env={'PEX_VERBOSE': '1'})
+    assert rc == 0, stdout
+
+
+def test_pex_script_function_custom_setuptools_useable():
+  with pex_with_entrypoints('my_app_function') as pex:
+    stdout, rc = run_simple_pex(pex, env={'PEX_VERBOSE': '1'})
+    assert rc == 0, stdout
+
+
+@contextmanager
+def pex_with_no_entrypoints():
+  with temporary_dir() as out:
+    pex = os.path.join(out, 'pex.pex')
+    run_pex_command(['setuptools==36.2.7', '-o', pex])
+    test_script = b'from setuptools.sandbox import run_setup; print(str(run_setup))'
+    yield pex, test_script, out
+
+
+def test_pex_interpreter_execute_custom_setuptools_useable():
+  with pex_with_no_entrypoints() as (pex, test_script, out):
+    script = os.path.join(out, 'script.py')
+    with open(script, 'wb') as fp:
+      fp.write(test_script)
+    stdout, rc = run_simple_pex(pex, args=(script,), env={'PEX_VERBOSE': '1'})
+    assert rc == 0, stdout
+
+
+def test_pex_interpreter_interact_custom_setuptools_useable():
+  with pex_with_no_entrypoints() as (pex, test_script, _):
+    stdout, rc = run_simple_pex(pex,
+                                env={'PEX_VERBOSE': '1'},
+                                stdin=test_script)
+    assert rc == 0, stdout
diff --git a/tests/test_pex.py b/tests/test_pex.py
--- a/tests/test_pex.py
+++ b/tests/test_pex.py
@@ -17,6 +17,7 @@
 from pex.pex import PEX
 from pex.pex_builder import PEXBuilder
 from pex.pex_info import PexInfo
+from pex.resolver import resolve
 from pex.testing import (
     IS_PYPY,
     ensure_python_interpreter,
@@ -359,3 +360,16 @@ def test_execute_interpreter_stdin_program():
     assert 0 == process.returncode
     assert b'one two\n' == stdout
     assert b'' == stderr
+
+
+def test_pex_run_custom_setuptools_useable():
+  with temporary_dir() as resolve_cache:
+    dists = resolve(['setuptools==36.2.7'], cache=resolve_cache)
+    with temporary_dir() as temp_dir:
+      pex = write_simple_pex(
+        temp_dir,
+        'from setuptools.sandbox import run_setup',
+        dists=dists,
+      )
+      rc = PEX(pex.path()).run()
+      assert rc == 0