Include PYTHONPATH in --inherit-path logic. (#765)

Previously `PEX` would inherit `PYTHONPATH` by default despite the
default `--inherit-path` setting being false. This was buggy behavior
running against a primary Pex goal of providing an isolated execution
environment. Now, in order to use `PYTHONPATH` to ammend the `sys.path`
of a running `PEX` a `PEX_INHERIT_PATH` (`pex --inherit-path=...`) of
`prefer` or `fallback` must be specified.

Update corresponding settings documentation and fix `PEX` to consider
`PYTHONPATH` when adjusting sys.path.

Fixes #707

Author: jsirois

pex/bin/pex.py

@@ -284,12 +284,11 @@ def configure_clp_pex_options(parser):
       default='false',
       action='store',
       choices=['false', 'fallback', 'prefer'],
-      help='Inherit the contents of sys.path (including site-packages) running the pex. '
-           'Possible values: false (does not inherit sys.path), '
-           'fallback (inherits sys.path after packaged dependencies), '
-           'prefer (inherits sys.path before packaged dependencies), '
-           'No value (alias for prefer, for backwards compatibility). '
-           '[Default: %default]')
+      help='Inherit the contents of sys.path (including site-packages, user site-packages and '
+           'PYTHONPATH) running the pex. Possible values: false (does not inherit sys.path), '
+           'fallback (inherits sys.path after packaged dependencies), prefer (inherits sys.path '
+           'before packaged dependencies), No value (alias for prefer, for backwards '
+           'compatibility). [Default: %default]')
 
   group.add_option(
       '--compile', '--no-compile',

pex/pex.py

@@ -202,6 +202,23 @@ def minimum_sys_modules(cls, site_libs, modules=None):
 
     return new_modules
 
+  _PYTHONPATH = 'PYTHONPATH'
+  _STASHED_PYTHONPATH = '_PEX_PYTHONPATH'
+
+  @classmethod
+  def stash_pythonpath(cls):
+    pythonpath = os.environ.pop(cls._PYTHONPATH, None)
+    if pythonpath is not None:
+      os.environ[cls._STASHED_PYTHONPATH] = pythonpath
+    return pythonpath
+
+  @classmethod
+  def unstash_pythonpath(cls):
+    pythonpath = os.environ.pop(cls._STASHED_PYTHONPATH, None)
+    if pythonpath is not None:
+      os.environ[cls._PYTHONPATH] = pythonpath
+    return pythonpath
+
   @classmethod
   def minimum_sys_path(cls, site_libs, inherit_path):
     scrub_paths = OrderedSet()
@@ -229,6 +246,27 @@ def all_distribution_paths(path):
         TRACER.log('Scrubbing from site-packages: %s' % path)
 
     scrubbed_sys_path = list(OrderedSet(sys.path) - scrub_paths)
+
+    pythonpath = cls.unstash_pythonpath()
+    if pythonpath is not None:
+      original_pythonpath = pythonpath.split(os.pathsep)
+      user_pythonpath = list(OrderedSet(original_pythonpath) - set(sys.path))
+      if original_pythonpath == user_pythonpath:
+        TRACER.log('Unstashed PYTHONPATH of %s' % pythonpath, V=2)
+      else:
+        TRACER.log('Extracted user PYTHONPATH of %s from unstashed PYTHONPATH of %s'
+                   % (os.pathsep.join(user_pythonpath), pythonpath), V=2)
+
+      if inherit_path == 'false':
+        for path in user_pythonpath:
+          TRACER.log('Scrubbing user PYTHONPATH element: %s' % path)
+      elif inherit_path == 'prefer':
+        TRACER.log('Prepending user PYTHONPATH: %s' % os.pathsep.join(user_pythonpath))
+        scrubbed_sys_path = user_pythonpath + scrubbed_sys_path
+      elif inherit_path == 'fallback':
+        TRACER.log('Appending user PYTHONPATH: %s' % os.pathsep.join(user_pythonpath))
+        scrubbed_sys_path = scrubbed_sys_path + user_pythonpath
+
     scrub_from_importer_cache = filter(
       lambda key: any(key.startswith(path) for path in scrub_paths),
       sys.path_importer_cache.keys())

pex/pex_bootstrapper.py

@@ -128,6 +128,11 @@ def maybe_reexec_pex(compatibility_constraints=None):
     # We've already been here and selected an interpreter. Continue to execution.
     return
 
+  from . import pex
+  pythonpath = pex.PEX.stash_pythonpath()
+  if pythonpath is not None:
+    TRACER.log('Stashed PYTHONPATH of {}'.format(pythonpath), V=2)
+
   with TRACER.timed('Selecting runtime interpreter', V=3):
     if ENV.PEX_PYTHON and not ENV.PEX_PYTHON_PATH:
       # preserve PEX_PYTHON re-exec for backwards compatibility
@@ -147,25 +152,38 @@ def maybe_reexec_pex(compatibility_constraints=None):
       )
       target = _select_path_interpreter(path=ENV.PEX_PYTHON_PATH,
                                         compatibility_constraints=compatibility_constraints)
-    else:
-      TRACER.log('Using the current interpreter {} since no constraints have been specified.'
-                 .format(sys.executable), V=3)
+    elif pythonpath is None:
+      TRACER.log('Using the current interpreter {} since no constraints have been specified and '
+                 'PYTHONPATH is not set.'.format(sys.executable), V=3)
       return
+    else:
+      target = current_interpreter
 
   os.environ.pop('PEX_PYTHON', None)
   os.environ.pop('PEX_PYTHON_PATH', None)
 
-  if target == current_interpreter:
-    TRACER.log('Using the current interpreter {} since it matches constraints.'
-               .format(sys.executable))
+  if pythonpath is None and target == current_interpreter:
+    TRACER.log('Using the current interpreter {} since it matches constraints and '
+               'PYTHONPATH is not set.'.format(sys.executable))
     return
 
   target_binary = target.binary
   cmdline = [target_binary] + sys.argv
-  TRACER.log('Re-executing: cmdline="%s", sys.executable="%s", PEX_PYTHON="%s", '
-             'PEX_PYTHON_PATH="%s", COMPATIBILITY_CONSTRAINTS="%s"'
-             % (cmdline, sys.executable, ENV.PEX_PYTHON, ENV.PEX_PYTHON_PATH,
-                compatibility_constraints))
+  TRACER.log(
+    'Re-executing: '
+    'cmdline={cmdline!r}, '
+    'sys.executable={python!r}, '
+    'PEX_PYTHON={pex_python!r}, '
+    'PEX_PYTHON_PATH={pex_python_path!r}, '
+    'COMPATIBILITY_CONSTRAINTS={compatibility_constraints!r}'
+    '{pythonpath}"'.format(
+      cmdline=' '.join(cmdline),
+      python=sys.executable,
+      pex_python=ENV.PEX_PYTHON,
+      pex_python_path=ENV.PEX_PYTHON_PATH,
+      compatibility_constraints=compatibility_constraints,
+      pythonpath=', (stashed) PYTHONPATH="{}"'.format(pythonpath) if pythonpath is not None else '')
+  )
 
   # Avoid a re-run through compatibility_constraint checking.
   os.environ[current_interpreter_blessed_env_var] = '1'

pex/pex_info.py

@@ -33,7 +33,8 @@ class PexInfo(object):
   script: string                      # script to execute in this pex environment
                                       # at most one of script/entry_point can be specified
   zip_safe: True, default False       # is this pex zip safe?
-  inherit_path: false/fallback/prefer # should this pex inherit site-packages + PYTHONPATH?
+  inherit_path: false/fallback/prefer # should this pex inherit site-packages + user site-packages
+                                      # + PYTHONPATH?
   ignore_errors: True, default False  # should we ignore inability to resolve dependencies?
   always_write_cache: False           # should we always write the internal cache to disk first?
                                       # this is useful if you have very large dependencies that

pex/testing.py

@@ -217,7 +217,8 @@ def write_simple_pex(td,
                      dists=None,
                      sources=None,
                      coverage=False,
-                     interpreter=None):
+                     interpreter=None,
+                     pex_info=None):
   """Write a pex file that optionally contains an executable entry point.
 
   :param str td: temporary directory path
@@ -229,6 +230,8 @@ def write_simple_pex(td,
   :param bool coverage: include coverage header
   :param interpreter: a custom interpreter to use to build the pex
   :type interpreter: :class:`pex.interpreter.PythonInterpreter`
+  :param pex_info: a custom PexInfo to use to build the pex.
+  :type pex_info: :class:`pex.pex_info.PexInfo`
   """
   dists = dists or []
   sources = sources or []
@@ -237,7 +240,8 @@ def write_simple_pex(td,
 
   pb = PEXBuilder(path=td,
                   preamble=COVERAGE_PREAMBLE if coverage else None,
-                  interpreter=interpreter)
+                  interpreter=interpreter,
+                  pex_info=pex_info)
 
   for dist in dists:
     pb.add_dist_location(dist.location)

pex/variables.py

@@ -176,12 +176,17 @@ def PEX_IGNORE_ERRORS(self):
 
   @property
   def PEX_INHERIT_PATH(self):
-    """Boolean
+    """String (false|prefer|fallback)
+
+    Allow inheriting packages from site-packages, user site-packages and the PYTHONPATH. By default,
+    PEX scrubs any non stdlib packages from sys.path prior to invoking the application. Using
+    'prefer' causes PEX to shift any non-stdlib packages before the pex environment on sys.path and
+    using 'fallback' shifts them after instead.
 
-    Allow inheriting packages from site-packages.  By default, PEX scrubs any packages and
-    namespace packages from sys.path prior to invoking the application.  This is generally not
-    advised, but can be used in situations when certain dependencies do not conform to standard
-    packaging practices and thus cannot be bundled into PEX files.  Default: false.
+    Using this option is generally not advised, but can help in situations when certain dependencies
+    do not conform to standard packaging practices and thus cannot be bundled into PEX files.
+
+    Default: false.
     """
     return self._get_string('PEX_INHERIT_PATH', default='false')
 

tests/test_integration.py

@@ -1525,7 +1525,8 @@ def test_issues_736_requirement_setup_py_with_extras():
 def _assert_exec_chain(exec_chain=None,
                        pex_python=None,
                        pex_python_path=None,
-                       interpreter_constraints=None):
+                       interpreter_constraints=None,
+                       pythonpath=None):
 
   with temporary_dir() as td:
     test_pex = os.path.join(td, 'test.pex')
@@ -1555,7 +1556,8 @@ def add_to_path(entry):
     env = make_env(_PEX_EXEC_CHAIN=1,
                    PEX_INTERPRETER=1,
                    PEX_PYTHON=pex_python,
-                   PEX_PYTHON_PATH=os.pathsep.join(pex_python_path) if pex_python_path else None)
+                   PEX_PYTHON_PATH=os.pathsep.join(pex_python_path) if pex_python_path else None,
+                   PYTHONPATH=os.pathsep.join(pythonpath) if pythonpath else None)
 
     output = subprocess.check_output(
       [sys.executable, test_pex, '-c', 'import json, os; print(json.dumps(os.environ.copy()))'],
@@ -1575,11 +1577,23 @@ def test_pex_no_reexec_no_constraints():
   _assert_exec_chain()
 
 
+def test_pex_reexec_no_constraints_pythonpath_present():
+  _assert_exec_chain(exec_chain=[os.path.realpath(sys.executable)],
+                     pythonpath=['.'])
+
+
 def test_pex_no_reexec_constraints_match_current():
   current_version = '.'.join(str(component) for component in sys.version_info[0:3])
   _assert_exec_chain(interpreter_constraints=['=={}'.format(current_version)])
 
 
+def test_pex_reexec_constraints_match_current_pythonpath_present():
+  current_version = '.'.join(str(component) for component in sys.version_info[0:3])
+  _assert_exec_chain(exec_chain=[os.path.realpath(sys.executable)],
+                     pythonpath=['.'],
+                     interpreter_constraints=['=={}'.format(current_version)])
+
+
 def test_pex_reexec_constraints_dont_match_current_pex_python_path():
   py36_interpreter = ensure_python_interpreter(PY36)
   py27_interpreter = ensure_python_interpreter(PY27)

tests/test_pex.py

@@ -6,11 +6,13 @@
 import sys
 import tempfile
 import textwrap
+from collections import namedtuple
 from contextlib import contextmanager
 from types import ModuleType
 
 import pytest
 
+from pex.common import safe_open
 from pex.compatibility import PY2, WINDOWS, nested, to_bytes
 from pex.installer import EggInstaller, WheelInstaller
 from pex.interpreter import PythonInterpreter
@@ -257,6 +259,105 @@ def test_pex_run():
       assert fake_stdout.read() == b'hellohello'
 
 
+class PythonpathIsolationTest(namedtuple('PythonpathIsolationTest',
+                                         ['pythonpath', 'dists', 'exe'])):
+
+  @staticmethod
+  def pex_info(inherit_path):
+    pex_info = PexInfo.default()
+    pex_info.inherit_path = inherit_path
+    return pex_info
+
+  def assert_isolation(self, inherit_path, expected_output):
+    env = dict(PYTHONPATH=self.pythonpath)
+    with named_temporary_file() as fake_stdout:
+      with temporary_dir() as temp_dir:
+        pex_builder = write_simple_pex(
+          temp_dir,
+          pex_info=self.pex_info(inherit_path),
+          dists=self.dists,
+          exe_contents=self.exe,
+        )
+
+        # Test the PEX.run API.
+        rc = PEX(pex_builder.path()).run(stdout=fake_stdout, env=env)
+        assert rc == 0
+
+        fake_stdout.seek(0)
+        assert expected_output == fake_stdout.read().decode('utf-8')
+
+        # Test direct PEX execution.
+        assert expected_output == subprocess.check_output([sys.executable, pex_builder.path()],
+                                                          env=env).decode('utf-8')
+
+
+@pytest.fixture(scope="module")
+def pythonpath_isolation_test():
+  with temporary_dir() as temp_dir:
+    pythonpath = os.path.join(temp_dir, 'one')
+    with safe_open(os.path.join(pythonpath, 'foo.py'), 'w') as fp:
+      fp.write('BAR = 42')
+    with safe_open(os.path.join(pythonpath, 'bar.py'), 'w') as fp:
+      fp.write('FOO = 137')
+
+    dist_content = {
+      'setup.py': textwrap.dedent("""
+        from setuptools import setup
+
+        setup(
+          name='foo',
+          version='0.0.0',
+          zip_safe=True,
+          packages=['foo'],
+          install_requires=[],
+        )
+      """),
+      'foo/__init__.py': 'BAR = 137',
+    }
+
+    with temporary_content(dist_content) as project_dir:
+      installer = WheelInstaller(project_dir)
+      foo_bdist = DistributionHelper.distribution_from_path(installer.bdist())
+
+      exe_contents = textwrap.dedent("""
+        import sys
+        
+        try:
+          import bar
+        except ImportError:
+          import collections
+          bar = collections.namedtuple('bar', ['FOO'])(None)
+  
+        import foo
+                  
+        sys.stdout.write("foo.BAR={} bar.FOO={}".format(foo.BAR, bar.FOO))
+      """)
+
+      yield PythonpathIsolationTest(pythonpath=pythonpath, dists=[foo_bdist], exe=exe_contents)
+
+
+def test_pythonpath_isolation_inherit_path_false(pythonpath_isolation_test):
+  pythonpath_isolation_test.assert_isolation(inherit_path='false',
+                                             expected_output='foo.BAR=137 bar.FOO=None')
+  # False should map to 'false'.
+  pythonpath_isolation_test.assert_isolation(inherit_path=False,
+                                             expected_output='foo.BAR=137 bar.FOO=None')
+
+
+def test_pythonpath_isolation_inherit_path_fallback(pythonpath_isolation_test):
+  pythonpath_isolation_test.assert_isolation(inherit_path='fallback',
+                                             expected_output='foo.BAR=137 bar.FOO=137')
+
+
+def test_pythonpath_isolation_inherit_path_prefer(pythonpath_isolation_test):
+  pythonpath_isolation_test.assert_isolation(inherit_path='prefer',
+                                             expected_output='foo.BAR=42 bar.FOO=137')
+
+  # True should map to 'prefer'.
+  pythonpath_isolation_test.assert_isolation(inherit_path=True,
+                                             expected_output='foo.BAR=42 bar.FOO=137')
+
+
 def test_pex_executable():
   # Tests that pex keeps executable permissions
   with temporary_dir() as temp_dir: