Repository: pantsbuild/pex
Issue: PEX_SCRIPT does not work with console_scripts

Due to an oversight, it only works with --script.  PEX_SCRIPT only checks for scripts within the Distribution data/metadata dirs.


Changed files:
pex/finders.py, pex/pex.py, pex/pex_builder.py

diff --git a/pex/finders.py b/pex/finders.py
index 1b8b0ef..37ddc12 100644
--- a/pex/finders.py
+++ b/pex/finders.py
@@ -275,3 +275,20 @@ def get_script_from_distributions(name, dists):
     if script_path:
       return dist, script_path, script_content
   return None, None, None
+
+
+def get_entry_point_from_console_script(script, dists):
+  # check all distributions for the console_script "script"
+  entries = frozenset(filter(None, (
+      dist.get_entry_map().get('console_scripts', {}).get(script) for dist in dists)))
+
+  # if multiple matches, freak out
+  if len(entries) > 1:
+    raise RuntimeError(
+        'Ambiguous script specification %s matches multiple entry points:%s' % (
+            script, ' '.join(map(str, entries))))
+
+  if entries:
+    entry_point = next(iter(entries))
+    # entry points are of the form 'foo = bar', we just want the 'bar' part:
+    return str(entry_point).split('=')[1].strip()
diff --git a/pex/pex.py b/pex/pex.py
index 94aa148..7f1acac 100644
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -16,7 +16,7 @@ from pkg_resources import EntryPoint, find_distributions
 from .common import die, safe_mkdir
 from .compatibility import exec_function
 from .environment import PEXEnvironment
-from .finders import get_script_from_distributions
+from .finders import get_entry_point_from_console_script, get_script_from_distributions
 from .interpreter import PythonInterpreter
 from .orderedset import OrderedSet
 from .pex_info import PexInfo
@@ -49,14 +49,13 @@ class PEX(object):  # noqa: T000
       sys.stderr.write('Could not bootstrap coverage module!\\n')
 
   @classmethod
-  def clean_environment(cls, forking=False):
+  def clean_environment(cls):
     try:
       del os.environ['MACOSX_DEPLOYMENT_TARGET']
     except KeyError:
       pass
-    if not forking:
-      for key in filter(lambda key: key.startswith('PEX_'), os.environ):
-        del os.environ[key]
+    for key in filter(lambda key: key.startswith('PEX_'), os.environ):
+      del os.environ[key]
 
   def __init__(self, pex=sys.argv[0], interpreter=None):
     self._pex = pex
@@ -306,14 +305,24 @@ class PEX(object):  # noqa: T000
       code.interact()
 
   def execute_script(self, script_name):
+    # TODO(wickman) This should be acheived by running clean_environment
+    # prior to invocation.
+    if 'PEX_SCRIPT' in os.environ:
+      del os.environ['PEX_SCRIPT']
+
     # TODO(wickman) PEXEnvironment should probably have a working_set property
     # or possibly just __iter__.
-    dist, script_path, script_content = get_script_from_distributions(
-        script_name, self._env.activate())
+    dists = list(self._env.activate())
+
+    entry_point = get_entry_point_from_console_script(script_name, dists)
+    if entry_point:
+      return 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))
-    self.execute_content(script_path, script_content, argv0=script_name)
+    return self.execute_content(script_path, script_content, argv0=script_name)
 
   @classmethod
   def execute_content(cls, name, content, argv0=None):
@@ -401,7 +410,7 @@ class PEX(object):  # noqa: T000
 
     Remaining keyword arguments are passed directly to subprocess.Popen.
     """
-    self.clean_environment(forking=True)
+    self.clean_environment()
 
     cmdline = self.cmdline(args)
     TRACER.log('PEX.run invoking %s' % ' '.join(cmdline))
diff --git a/pex/pex_builder.py b/pex/pex_builder.py
index 67785f8..c2419ed 100644
--- a/pex/pex_builder.py
+++ b/pex/pex_builder.py
@@ -10,7 +10,7 @@ from pkg_resources import DefaultProvider, ZipProvider, get_provider
 
 from .common import Chroot, chmod_plus_x, open_zip, safe_mkdir, safe_mkdtemp
 from .compatibility import to_bytes
-from .finders import get_script_from_distributions
+from .finders import get_entry_point_from_console_script, get_script_from_distributions
 from .interpreter import PythonInterpreter
 from .marshaller import CodeMarshaller
 from .pex_info import PexInfo
@@ -183,23 +183,6 @@ class PEXBuilder(object):
     entry_point.replace(os.path.sep, '.')
     self._pex_info.entry_point = entry_point.rpartition('.')[0]
 
-  @classmethod
-  def _get_entry_point_from_console_script(cls, script, distributions):
-    # check all distributions for the console_script "script"
-    entries = frozenset(filter(None, (
-        dist.get_entry_map().get('console_scripts', {}).get(script) for dist in distributions)))
-
-    # if multiple matches, freak out
-    if len(entries) > 1:
-      raise cls.InvalidExecutableSpecification(
-          'Ambiguous script specification %s matches multiple entry points:%s' % (
-              script, ' '.join(map(str, entries))))
-
-    if entries:
-      entry_point = next(iter(entries))
-      # entry points are of the form 'foo = bar', we just want the 'bar' part:
-      return str(entry_point).split('=')[1].strip()
-
   def set_script(self, script):
     """Set the entry point of this PEX environment based upon a distribution script.
 
@@ -210,7 +193,7 @@ class PEXBuilder(object):
     """
 
     # check if 'script' is a console_script
-    entry_point = self._get_entry_point_from_console_script(script, self._distributions)
+    entry_point = get_entry_point_from_console_script(script, self._distributions)
     if entry_point:
       self.set_entry_point(entry_point)
       return