Improve wheel support in pex.

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()