Treat .pth injected paths as extras. (#370)

Fixes #302

before:

    [omerta pex (kwlzn/osxtras)]$ pex --version
    pex 1.2.3
    [omerta pex (kwlzn/osxtras)]$ pex six -o /tmp/six_broke.pex
    [omerta pex (kwlzn/osxtras)]$ PEX_VERBOSE=9 PEX_PYTHON=/usr/bin/python /tmp/six_broke.pex 
    pex: Detected PEX_PYTHON, re-exec to /usr/bin/python
    ...
    Python 2.7.10 (default, Jul 30 2016, 19:40:32) 
    [GCC 4.2.1 Compatible Apple LLVM 8.0.0 (clang-800.0.34)] on darwin
    Type "help", "copyright", "credits" or "license" for more information.
    (InteractiveConsole)
    >>> from six import wraps
    Traceback (most recent call last):
      File "<console>", line 1, in <module>
    ImportError: cannot import name wraps
    >>> import six; six.__file__
    '/System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python/six.pyc'
    >>>

after:

    [omerta pex (kwlzn/osxtras)]$ tox -e py27-package
    ...
    [omerta pex (kwlzn/osxtras)]$ ./dist/pex27 six -o /tmp/six.pex
    ...
    [omerta pex (kwlzn/osxtras)]$ PEX_VERBOSE=9 PEX_PYTHON=/usr/bin/python /tmp/six.pex 
    pex: Detected PEX_PYTHON, re-exec to /usr/bin/python
    ...
    pex: Found .pth file: /Library/Python/2.7/site-packages/Extras.pth
    pex: Found site extra: /System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python
    pex: Found site extra: /System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python/PyObjC
    ...
    pex: Tainted path element: /System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python
    pex: Tainted path element: /System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python/PyObjC
    ...
    pex: Scrubbing from site-packages: /System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python
    pex: Scrubbing from site-packages: /System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python/PyObjC
    ...
    Python 2.7.10 (default, Jul 30 2016, 19:40:32) 
    [GCC 4.2.1 Compatible Apple LLVM 8.0.0 (clang-800.0.34)] on darwin
    Type "help", "copyright", "credits" or "license" for more information.
    (InteractiveConsole)
    >>> from six import wraps
    >>> import six; six.__file__
    '/Users/kwilson/.pex/install/six-1.10.0-py2.py3-none-any.whl.a99dfb27e60da3957f6667853b91ea52e173da0a/six-1.10.0-py2.py3-none-any.whl/six.pyc'
    >>>

pex/pex.py

@@ -21,6 +21,7 @@
 from .orderedset import OrderedSet
 from .pex_info import PexInfo
 from .tracer import TRACER
+from .util import iter_pth_paths
 from .variables import ENV
 
 
@@ -88,18 +89,38 @@ def _activate(self):
   @classmethod
   def _extras_paths(cls):
     standard_lib = sysconfig.get_python_lib(standard_lib=True)
+
     try:
       makefile = sysconfig.parse_makefile(sysconfig.get_makefile_filename())
     except (AttributeError, IOError):
       # This is not available by default in PyPy's distutils.sysconfig or it simply is
       # no longer available on the system (IOError ENOENT)
       makefile = {}
+
     extras_paths = filter(None, makefile.get('EXTRASPATH', '').split(':'))
     for path in extras_paths:
       yield os.path.join(standard_lib, path)
 
-  @classmethod
-  def _get_site_packages(cls):
+    # Handle .pth injected paths as extras.
+    sitedirs = cls._get_site_packages()
+    for pth_path in cls._scan_pth_files(sitedirs):
+      TRACER.log('Found .pth file: %s' % pth_path, V=3)
+      for extras_path in iter_pth_paths(pth_path):
+        yield extras_path
+
+  @staticmethod
+  def _scan_pth_files(dir_paths):
+    """Given an iterable of directory paths, yield paths to all .pth files within."""
+    for dir_path in dir_paths:
+      if not os.path.exists(dir_path):
+        continue
+
+      pth_filenames = (f for f in os.listdir(dir_path) if f.endswith('.pth'))
+      for pth_filename in pth_filenames:
+        yield os.path.join(dir_path, pth_filename)
+
+  @staticmethod
+  def _get_site_packages():
     try:
       from site import getsitepackages
       return set(getsitepackages())

pex/util.py

@@ -9,11 +9,13 @@
 import tempfile
 import uuid
 from hashlib import sha1
+from site import makepath
 from threading import Lock
 
 from pkg_resources import find_distributions, resource_isdir, resource_listdir, resource_string
 
 from .common import rename_if_empty, safe_mkdir, safe_mkdtemp, safe_open
+from .compatibility import exec_function
 from .finders import register_finders
 
 
@@ -213,3 +215,33 @@ def named_temporary_file(*args, **kwargs):
       yield fp
   finally:
     os.remove(fp.name)
+
+
+def iter_pth_paths(filename):
+  """Given a .pth file, extract and yield all inner paths without honoring imports. This shadows
+  python's site.py behavior, which is invoked at interpreter startup."""
+  try:
+    f = open(filename, 'rU')  # noqa
+  except IOError:
+    return
+
+  dirname = os.path.dirname(filename)
+  known_paths = set()
+
+  with f:
+    for line in f:
+      line = line.rstrip()
+      if not line or line.startswith('#'):
+        continue
+      elif line.startswith(('import ', 'import\t')):
+        try:
+          exec_function(line)
+          continue
+        except Exception:
+          # Defer error handling to the higher level site.py logic invoked at startup.
+          return
+      else:
+        extras_dir, extras_dir_case_insensitive = makepath(dirname, line)
+        if extras_dir_case_insensitive not in known_paths and os.path.exists(extras_dir):
+          yield extras_dir
+          known_paths.add(extras_dir_case_insensitive)

tests/test_util.py

@@ -11,11 +11,11 @@
 from twitter.common.contextutil import temporary_dir
 
 from pex.common import safe_mkdir
-from pex.compatibility import nested
+from pex.compatibility import nested, to_bytes
 from pex.installer import EggInstaller, WheelInstaller
 from pex.pex_builder import PEXBuilder
 from pex.testing import make_bdist, run_simple_pex, temporary_content, write_zipfile
-from pex.util import CacheHelper, DistributionHelper, named_temporary_file
+from pex.util import CacheHelper, DistributionHelper, iter_pth_paths, named_temporary_file
 
 try:
   from unittest import mock
@@ -174,3 +174,35 @@ def test_distributionhelper_egg_assert():
     'setuptools'
   )
   assert len(d.resource_listdir('/')) > 3
+
+
+@mock.patch('os.path.exists', autospec=True, spec_set=True)
+def test_iter_pth_paths(mock_exists):
+  # Ensure path checking always returns True for dummy paths.
+  mock_exists.return_value = True
+
+  with temporary_dir() as tmpdir:
+    in_tmp = lambda f: os.path.join(tmpdir, f)
+
+    PTH_TEST_MAPPING = {
+      # A mapping of .pth file content -> expected paths.
+      '/System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python\n': [
+        '/System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python'
+      ],
+      'relative_path\nrelative_path2\n\nrelative_path3': [
+        in_tmp('relative_path'),
+        in_tmp('relative_path2'),
+        in_tmp('relative_path3')
+      ],
+      'duplicate_path\nduplicate_path': [in_tmp('duplicate_path')],
+      'randompath\nimport nosuchmodule\n': [in_tmp('randompath')],
+      'import nosuchmodule\nfoo': [],
+      'import nosuchmodule\n': [],
+      'import bad)syntax\n': [],
+    }
+
+    for i, pth_content in enumerate(PTH_TEST_MAPPING):
+      pth_tmp_path = os.path.abspath(os.path.join(tmpdir, 'test%s.pth' % i))
+      with open(pth_tmp_path, 'wb') as f:
+        f.write(to_bytes(pth_content))
+      assert sorted(PTH_TEST_MAPPING[pth_content]) == sorted(list(iter_pth_paths(pth_tmp_path)))