Improve wheel support in pex. (#388)

Fixes #325

Wheel handling in pex files is broken. The wheel standard says
that wheel files are not designed to be importable archives,
but pex treats them as if they are. This causes many standard
compliant wheels, including things like tensorflow and opencv,
to fail to import in pexes (and thus in pants).

This change modifies wheel handling, so that when a wheel
is added to a pex, it's installed in an importable form.

pex/finders.py

@@ -245,19 +245,22 @@ def get_script_from_egg(name, dist):
   return None, None
 
 
-def safer_name(name):
-  return name.replace('-', '_')
-
-
 def get_script_from_whl(name, dist):
-  # This is true as of at least wheel==0.24.  Might need to take into account the
-  # metadata version bundled with the wheel.
-  wheel_scripts_dir = '%s-%s.data/scripts' % (safer_name(dist.key), dist.version)
-  if dist.resource_isdir(wheel_scripts_dir) and name in dist.resource_listdir(wheel_scripts_dir):
-    script_path = os.path.join(wheel_scripts_dir, name)
-    return (
-        os.path.join(dist.egg_info, script_path),
-        dist.get_resource_string('', script_path).replace(b'\r\n', b'\n').replace(b'\r', b'\n'))
+  # This can get called in different contexts; in some, it looks for files in the
+  # wheel archives being used to produce a pex; in others, it looks for files in the
+  # install wheel directory included in the pex. So we need to look at both locations.
+  datadir_name = "%s-%s.data" % (dist.project_name, dist.version)
+  wheel_scripts_dirs = ['bin', 'scripts',
+                         os.path.join(datadir_name, "bin"),
+                         os.path.join(datadir_name, "scripts")]
+  for wheel_scripts_dir in wheel_scripts_dirs:
+    if (dist.resource_isdir(wheel_scripts_dir) and
+        name in dist.resource_listdir(wheel_scripts_dir)):
+      # We always install wheel scripts into bin
+      script_path = os.path.join(wheel_scripts_dir, name)
+      return (
+          os.path.join(dist.egg_info, script_path),
+          dist.get_resource_string('', script_path).replace(b'\r\n', b'\n').replace(b'\r', b'\n'))
   return None, None
 
 

pex/interpreter.py

@@ -400,6 +400,10 @@ def get_location(self, req):
       if req.key == dist_name and dist_version in req:
         return location
 
+  def supports_wheel_install(self):
+    """Wheel installs are broken in python 2.6"""
+    return self.version >= (2, 7)
+
   def __hash__(self):
     return hash((self._binary, self._identity))
 

pex/pex_builder.py

@@ -1,7 +1,7 @@
 # Copyright 2014 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-from __future__ import absolute_import
+from __future__ import absolute_import, print_function
 
 import logging
 import os
@@ -253,7 +253,45 @@ def _add_dist_dir(self, path, dist_name):
         self._copy_or_link(filename, target)
     return CacheHelper.dir_hash(path)
 
+  def _get_installer_paths(self, base):
+    """Set up an overrides dict for WheelFile.install that installs the contents
+    of a wheel into its own base in the pex dependencies cache.
+    """
+    return {
+      'purelib': base,
+      'headers': os.path.join(base, 'headers'),
+      'scripts': os.path.join(base, 'bin'),
+      'platlib': base,
+      'data': base
+    }
+
   def _add_dist_zip(self, path, dist_name):
+    # We need to distinguish between wheels and other zips. Most of the time,
+    # when we have a zip, it contains its contents in an importable form.
+    # But wheels don't have to be importable, so we need to force them
+    # into an importable shape. We can do that by installing it into its own
+    # wheel dir.
+    if not self.interpreter.supports_wheel_install():
+      self._logger.warn("Wheel dependency on %s may not work correctly with Python 2.6." %
+                        dist_name)
+
+    if self.interpreter.supports_wheel_install() and dist_name.endswith("whl"):
+      from wheel.install import WheelFile
+      tmp = safe_mkdtemp()
+      whltmp = os.path.join(tmp, dist_name)
+      os.mkdir(whltmp)
+      wf = WheelFile(path)
+      wf.install(overrides=self._get_installer_paths(whltmp), force=True)
+      for (root, _, files) in os.walk(whltmp):
+        pruned_dir = os.path.relpath(root, tmp)
+        for f in files:
+          fullpath = os.path.join(root, f)
+          if os.path.isdir(fullpath):
+            continue
+          target = os.path.join(self._pex_info.internal_cache, pruned_dir, f)
+          self._chroot.copy(fullpath, target)
+      return CacheHelper.dir_hash(whltmp)
+
     with open_zip(path) as zf:
       for name in zf.namelist():
         if name.endswith('/'):

setup.py

@@ -54,6 +54,7 @@
   ],
   install_requires = [
     SETUPTOOLS_REQUIREMENT,
+    WHEEL_REQUIREMENT,
   ],
   tests_require = [
     'mock',

tests/test_compiler.py

@@ -2,9 +2,9 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import contextlib
+import marshal
 import os
 
-import marshal
 import pytest
 from twitter.common.contextutil import temporary_dir
 

tests/test_pex.py

@@ -186,6 +186,8 @@ def test_site_libs_excludes_prefix():
     assert sys.prefix not in site_libs
 
 
+@pytest.mark.skipif(sys.version_info < (2, 7),
+                    reason="wheel script installation is broken on python 2.6")
 @pytest.mark.parametrize('zip_safe', (False, True))
 @pytest.mark.parametrize('project_name', ('my_project', 'my-project'))
 @pytest.mark.parametrize('installer_impl', (EggInstaller, WheelInstaller))

tests/test_pex_builder.py

@@ -3,6 +3,7 @@
 
 import os
 import stat
+import sys
 import zipfile
 from contextlib import closing
 
@@ -27,6 +28,16 @@
   fp.write('success')
 """
 
+wheeldeps_exe_main = """
+import sys
+from pyparsing import *
+from my_package.my_module import do_something
+do_something()
+
+with open(sys.argv[1], 'w') as fp:
+  fp.write('success')
+"""
+
 
 def test_pex_builder():
   # test w/ and w/o zipfile dists
@@ -57,6 +68,23 @@ def test_pex_builder():
       assert fp.read() == 'success'
 
 
+@pytest.mark.skipif(sys.version_info < (2, 7),
+                    reason="wheel script installation is broken on python 2.6")
+def test_pex_builder_wheeldep():
+  """Repeat the pex_builder test, but this time include an import of
+  something from a wheel that doesn't come in importable form.
+  """
+  with nested(temporary_dir(), make_bdist('p1', zipped=True)) as (td, p1):
+    pyparsing_path = "./tests/example_packages/pyparsing-2.1.10-py2.py3-none-any.whl"
+    dist = DistributionHelper.distribution_from_path(pyparsing_path)
+    write_pex(td, wheeldeps_exe_main, dists=[p1, dist])
+    success_txt = os.path.join(td, 'success.txt')
+    PEX(td).run(args=[success_txt])
+    assert os.path.exists(success_txt)
+    with open(success_txt) as fp:
+      assert fp.read() == 'success'
+
+
 def test_pex_builder_shebang():
   def builder(shebang):
     pb = PEXBuilder()