Fix script search in wheels.

We had approximately the right handling of wheel name components in the
same file but were not leveraging common code. Robustify name component
handling, leverage and add tests.

Fixes pex-tool#443
Fixes pex-tool#551

pex/finders.py

@@ -16,6 +16,7 @@
 
 import os
 import pkgutil
+import re
 import sys
 import zipimport
 
@@ -99,12 +100,39 @@ def _remove_finder(importer, finder):
 class WheelMetadata(pkg_resources.EggMetadata):
   """Metadata provider for zipped wheels."""
 
+  @classmethod
+  def _escape(cls, filename_component):
+    # See: https://www.python.org/dev/peps/pep-0427/#escaping-and-unicode
+    return re.sub("[^\w\d.]+", "_", filename_component, re.UNICODE)
+
   @classmethod
   def _split_wheelname(cls, wheelname):
-    split_wheelname = wheelname.rsplit('-', 4)
-    assert len(split_wheelname) == 5, 'invalid wheel name: %s' % (wheelname)
-    split_wheelname[0] = split_wheelname[0].replace('-', '_')
-    return '-'.join(split_wheelname[:-3])
+    # See: https://www.python.org/dev/peps/pep-0427/#file-name-convention
+    assert wheelname.endswith('.whl'), 'invalid wheel name: %s' % wheelname
+    split_wheelname = wheelname.rsplit('-', 5)
+    assert len(split_wheelname) in (5, 6), 'invalid wheel name: %s' % wheelname
+    distribution, version = split_wheelname[:2]
+    return '%s-%s' % (distribution, version)
+
+  @classmethod
+  def data_dir(cls, wheel_path):
+    """Returns the internal path of the data dir for the given wheel.
+
+    As defined https://www.python.org/dev/peps/pep-0427/#the-data-directory
+
+    :rtype: str
+    """
+    return '%s.data' % cls._split_wheelname(os.path.basename(wheel_path))
+
+  @classmethod
+  def dist_info_dir(cls, wheel_path):
+    """Returns the internal path of the dist-info dir for the given wheel.
+
+    As defined here: https://www.python.org/dev/peps/pep-0427/#the-dist-info-directory
+
+    :rtype: str
+    """
+    return '%s.dist-info' % cls._split_wheelname(os.path.basename(wheel_path))
 
   def _setup_prefix(self):
     path = self.module_path
@@ -114,7 +142,7 @@ def _setup_prefix(self):
         self.egg_name = os.path.basename(path)
         # TODO(wickman) Test the regression where we have both upper and lower cased package
         # names.
-        self.egg_info = os.path.join(path, '%s.dist-info' % self._split_wheelname(self.egg_name))
+        self.egg_info = os.path.join(path, self.dist_info_dir(self.egg_name))
         self.egg_root = path
         break
       old = path
@@ -251,7 +279,7 @@ def get_script_from_whl(name, dist):
   # 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)
+  datadir_name = WheelMetadata.data_dir(dist.location)
   wheel_scripts_dirs = ['bin', 'scripts',
                          os.path.join(datadir_name, "bin"),
                          os.path.join(datadir_name, "scripts")]
@@ -261,7 +289,7 @@ def get_script_from_whl(name, dist):
       # 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),
+          os.path.join(dist.location, script_path),
           dist.get_resource_string('', script_path).replace(b'\r\n', b'\n').replace(b'\r', b'\n'))
   return None, None
 
@@ -300,15 +328,22 @@ def get_entrypoint(dist):
     script_entry = dist.get_entry_map().get('console_scripts', {}).get(script)
     if script_entry is not None:
       # Entry points are of the form 'foo = bar', we just want the 'bar' part.
-      return dist.key, str(script_entry).split('=')[1].strip()
+      return str(script_entry).split('=')[1].strip()
 
-  entries = frozenset(filter(None, (get_entrypoint(dist) for dist in dists)))
+  entries = {}
+  for dist in dists:
+    entry_point = get_entrypoint(dist)
+    if entry_point is not None:
+      entries[dist.key] = (dist, entry_point)
 
   if len(entries) > 1:
     raise RuntimeError(
         'Ambiguous script specification %s matches multiple entry points:\n\t%s' % (
-            script, '\n\t'.join('%s from %s' % (entry_point, key) for key, entry_point in entries)))
+            script,
+            '\n\t'.join('%r from %r' % (entry_point, dist)
+                        for dist, entry_point in entries.values())))
 
+  dist, entry_point = None, None
   if entries:
-    _, entry_point = next(iter(entries))
-    return entry_point
+    dist, entry_point = next(iter(entries.values()))
+  return dist, entry_point

pex/pex.py

@@ -463,14 +463,15 @@ def execute_interpreter(self):
   def execute_script(self, script_name):
     dists = list(self._activate())
 
-    entry_point = get_entry_point_from_console_script(script_name, dists)
+    dist, entry_point = get_entry_point_from_console_script(script_name, dists)
     if entry_point:
+      TRACER.log('Found console_script %r in %r' % (entry_point, dist))
       sys.exit(self.execute_entry(entry_point))
 
     dist, script_path, script_content = get_script_from_distributions(script_name, dists)
     if not dist:
-      raise self.NotFound('Could not find script %s in pex!' % script_name)
-    TRACER.log('Found script %s in %s' % (script_name, dist))
+      raise self.NotFound('Could not find script %r in pex!' % script_name)
+    TRACER.log('Found script %r in %r' % (script_name, dist))
     return self.execute_content(script_path, script_content, argv0=script_name)
 
   @classmethod

pex/pex_builder.py

@@ -14,6 +14,7 @@
 from .finders import get_entry_point_from_console_script, get_script_from_distributions
 from .interpreter import PythonInterpreter
 from .pex_info import PexInfo
+from .tracer import TRACER
 from .util import CacheHelper, DistributionHelper
 
 BOOTSTRAP_ENVIRONMENT = b"""
@@ -207,17 +208,19 @@ def set_script(self, script):
     """
 
     # check if 'script' is a console_script
-    entry_point = get_entry_point_from_console_script(script, self._distributions)
+    dist, entry_point = get_entry_point_from_console_script(script, self._distributions)
     if entry_point:
       self.set_entry_point(entry_point)
+      TRACER.log('Set entrypoint to console_script %r in %r' % (entry_point, dist))
       return
 
     # check if 'script' is an ordinary script
-    script_path, _, _ = get_script_from_distributions(script, self._distributions)
-    if script_path:
+    dist, _, _ = get_script_from_distributions(script, self._distributions)
+    if dist:
       if self._pex_info.entry_point:
         raise self.InvalidExecutableSpecification('Cannot set both entry point and script of PEX!')
       self._pex_info.script = script
+      TRACER.log('Set entrypoint to script %r in %r' % (script, dist))
       return
 
     raise self.InvalidExecutableSpecification(

tests/test_finders.py

@@ -1,6 +1,7 @@
 # Copyright 2014 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+import os
 import zipimport
 
 import pkg_resources
@@ -9,7 +10,13 @@
 from pex.finders import ChainedFinder
 from pex.finders import _add_finder as add_finder
 from pex.finders import _remove_finder as remove_finder
-from pex.finders import find_eggs_in_zip, get_entry_point_from_console_script, get_script_from_egg
+from pex.finders import (
+    find_eggs_in_zip,
+    find_wheels_in_zip,
+    get_entry_point_from_console_script,
+    get_script_from_egg,
+    get_script_from_whl
+)
 
 try:
   import mock
@@ -105,20 +112,46 @@ def test_remove_finder():
       mock_register_finder.assert_called_with('foo', pkg_resources.find_nothing)
 
 
+def test_get_script_from_egg_with_no_scripts():
+  # Make sure eggs without scripts don't cause errors.
+  egg_path = './tests/example_packages/Flask_Cache-0.13.1-py2.7.egg'
+  dists = list(find_eggs_in_zip(zipimport.zipimporter(egg_path), egg_path, only=True))
+  assert len(dists) == 1
+
+  dist = dists[0]
+  assert (None, None) == get_script_from_egg('non_existent_script', dist)
+
+
 def test_get_script_from_egg():
-  # Make sure eggs without scripts don't cause errors
-  dists = list(
-    find_eggs_in_zip(
-      zipimport.zipimporter('./tests/example_packages/Flask_Cache-0.13.1-py2.7.egg'),
-      './tests/example_packages/Flask_Cache-0.13.1-py2.7.egg',
-      only=True))
+  egg_path = './tests/example_packages/eno-0.0.17-py2.7.egg'
+  dists = list(find_eggs_in_zip(zipimport.zipimporter(egg_path), egg_path, only=True))
+  assert len(dists) == 1
+
+  dist = dists[0]
 
+  location, content = get_script_from_egg('run_eno_server', dist)
+  assert os.path.join(egg_path, 'EGG-INFO/scripts/run_eno_server') == location
+  assert content.startswith('#!'), 'Expected a `scripts` style script with shebang.'
+
+  assert (None, None) == get_script_from_egg('non_existent_script', dist)
+
+
+# In-part, tests a bug where the wheel distribution name has dashes as reported in:
+#   https://github.com/pantsbuild/pex/issues/443
+#   https://github.com/pantsbuild/pex/issues/551
+def test_get_script_from_whl():
+  whl_path = './tests/example_packages/aws_cfn_bootstrap-1.4-py2-none-any.whl'
+  dists = list(find_wheels_in_zip(zipimport.zipimporter(whl_path), whl_path))
   assert len(dists) == 1
 
-  (location, content) = get_script_from_egg('non_existent_script', dists[0])
+  dist = dists[0]
+  assert 'aws-cfn-bootstrap' == dist.project_name
 
-  assert location is None
-  assert content is None
+  script_path, script_content = get_script_from_whl('cfn-signal', dist)
+  assert os.path.join(whl_path, 'aws_cfn_bootstrap-1.4.data/scripts/cfn-signal') == script_path
+  assert script_content.startswith('#!'), 'Expected a `scripts` style script with shebang.'
+
+  assert (None, None) == get_script_from_whl('non_existent_script', dist)
 
 
 class FakeDist(object):
@@ -134,7 +167,10 @@ def get_entry_map(self):
 def test_get_entry_point_from_console_script():
   dists = [FakeDist(key='fake', console_script_entry='bob= bob.main:run'),
            FakeDist(key='fake', console_script_entry='bob =bob.main:run')]
-  assert 'bob.main:run' == get_entry_point_from_console_script('bob', dists)
+
+  dist, entrypoint = get_entry_point_from_console_script('bob', dists)
+  assert 'bob.main:run' == entrypoint
+  assert dist in dists
 
 
 def test_get_entry_point_from_console_script_conflict():
@@ -147,4 +183,4 @@ def test_get_entry_point_from_console_script_conflict():
 def test_get_entry_point_from_console_script_dne():
   dists = [FakeDist(key='bob', console_script_entry='bob= bob.main:run'),
            FakeDist(key='fake', console_script_entry='bob =bob.main:run')]
-  assert None is get_entry_point_from_console_script('jane', dists)
+  assert (None, None) == get_entry_point_from_console_script('jane', dists)

tests/test_pex.py

@@ -211,7 +211,7 @@ def test_pex_script(installer_impl, project_name, zip_safe):
     env_copy['PEX_SCRIPT'] = 'hello_world'
     so, rc = run_simple_pex_test('', env=env_copy)
     assert rc == 1, so.decode('utf-8')
-    assert b'Could not find script hello_world' in so
+    assert b"Could not find script 'hello_world'" in so
 
     so, rc = run_simple_pex_test('', env=env_copy, dists=[bdist])
     assert rc == 0, so.decode('utf-8')