Validate entry point at build time (pex-tool#521)

Resolves pex-tool#508

### Problem

Current behavior allows building a pex that cannot execute:

```
[omerta ~]$ pex requests -e qqqqqqqqqqqq -o test.pex
[omerta ~]$ ./test.pex
Traceback (most recent call last):
  File ".bootstrap/_pex/pex.py", line 367, in execute
  File ".bootstrap/_pex/pex.py", line 293, in _wrap_coverage
  File ".bootstrap/_pex/pex.py", line 325, in _wrap_profiling
  File ".bootstrap/_pex/pex.py", line 410, in _execute
  File ".bootstrap/_pex/pex.py", line 468, in execute_entry
  File ".bootstrap/_pex/pex.py", line 473, in execute_module
  File "/Users/kwilson/Python/CPython-2.7.13/lib/python2.7/runpy.py", line 182, in run_module
    mod_name, loader, code, fname = _get_module_details(mod_name)
  File "/Users/kwilson/Python/CPython-2.7.13/lib/python2.7/runpy.py", line 107, in _get_module_details
    raise error(format(e))
ImportError: No module named qqqqqqqqqqqq
```

### Solution

Verify entry point at build time. E.g. `a.b.c:m` means we will try to do `from a.b.c import m` in a separate process.

### Result

```
$ find hello
hello
hello/test.py
hello/tree.py

# Invalid module
$ python bin.py -D hello -e invalid.module  -o x.pex --validate-entry-point
Traceback (most recent call last):
  File "<string>", line 1, in <module>
ImportError: No module named invalid.module
Traceback (most recent call last):
  File "bin.py", line 758, in <module>
    main()
  File "bin.py", line 743, in main
    pex_builder.build(tmp_name, verify_entry_point=options.validate_ep)
  File "/Users/yic/workspace/pex/pex/pex_builder.py", line 529, in build
    self.freeze(bytecode_compile=bytecode_compile, verify_entry_point=verify_entry_point)
  File "/Users/yic/workspace/pex/pex/pex_builder.py", line 514, in freeze
    self._verify_entry_point()
  File "/Users/yic/workspace/pex/pex/pex_builder.py", line 263, in _verify_entry_point
    raise self.InvalidEntryPoint('Failed to:`{}`'.format(import_statement))
pex.pex_builder.InvalidEntryPoint: Failed to:`import invalid.module`

# invalid method
$ python bin.py -D hello -e test:invalid_method  -o x.pex --validate-entry-point
Traceback (most recent call last):
  File "<string>", line 1, in <module>
ImportError: cannot import name invalid_method
Traceback (most recent call last):
  File "bin.py", line 758, in <module>
    main()
  File "bin.py", line 743, in main
    pex_builder.build(tmp_name, verify_entry_point=options.validate_ep)
  File "/Users/yic/workspace/pex/pex/pex_builder.py", line 529, in build
    self.freeze(bytecode_compile=bytecode_compile, verify_entry_point=verify_entry_point)
  File "/Users/yic/workspace/pex/pex/pex_builder.py", line 514, in freeze
    self._verify_entry_point()
  File "/Users/yic/workspace/pex/pex/pex_builder.py", line 263, in _verify_entry_point
    raise self.InvalidEntryPoint('Failed to:`{}`'.format(import_statement))
pex.pex_builder.InvalidEntryPoint: Failed to:`from test import invalid_method`

# without the flag, invalid method still works
$ python bin.py -D hello -e test:invalid_method  -o x.pex 
```

pex/bin/pex.py

@@ -386,6 +386,15 @@ def configure_clp_pex_entry_points(parser):
       help='Set the entry point as to the script or console_script as defined by a any of the '
            'distributions in the pex.  For example: "pex -c fab fabric" or "pex -c mturk boto".')
 
+  group.add_option(
+      '--validate-entry-point',
+      dest='validate_ep',
+      default=False,
+      action='store_true',
+      help='Validate the entry point by importing it in separate process. Warning: this could have '
+           'side effects. For example, entry point `a.b.c:m` will translate to '
+           '`from a.b.c import m` during validation. [Default: %default]')
+
   parser.add_option_group(group)
 
 
@@ -727,22 +736,24 @@ def main(args=None):
     with TRACER.timed('Building pex'):
       pex_builder = build_pex(reqs, options, resolver_options_builder)
 
+    pex_builder.freeze()
+    pex = PEX(pex_builder.path(),
+              interpreter=pex_builder.interpreter,
+              verify_entry_point=options.validate_ep)
+
     if options.pex_name is not None:
       log('Saving PEX file to %s' % options.pex_name, v=options.verbosity)
       tmp_name = options.pex_name + '~'
       safe_delete(tmp_name)
       pex_builder.build(tmp_name)
       os.rename(tmp_name, options.pex_name)
-      return 0
-
-    if not _compatible_with_current_platform(options.platforms):
-      log('WARNING: attempting to run PEX with incompatible platforms!')
-
-    pex_builder.freeze()
+    else:
+      if not _compatible_with_current_platform(options.platforms):
+        log('WARNING: attempting to run PEX with incompatible platforms!')
 
-    log('Running PEX file at %s with args %s' % (pex_builder.path(), cmdline), v=options.verbosity)
-    pex = PEX(pex_builder.path(), interpreter=pex_builder.interpreter)
-    sys.exit(pex.run(args=list(cmdline)))
+      log('Running PEX file at %s with args %s' % (pex_builder.path(), cmdline),
+          v=options.verbosity)
+      sys.exit(pex.run(args=list(cmdline)))
 
 
 if __name__ == '__main__':

pex/pex.py

@@ -21,7 +21,7 @@
 from .orderedset import OrderedSet
 from .pex_info import PexInfo
 from .tracer import TRACER
-from .util import iter_pth_paths, merge_split
+from .util import iter_pth_paths, merge_split, named_temporary_file
 from .variables import ENV
 
 
@@ -41,6 +41,7 @@ class PEX(object):  # noqa: T000
 
   class Error(Exception): pass
   class NotFound(Error): pass
+  class InvalidEntryPoint(Error): pass
 
   @classmethod
   def clean_environment(cls):
@@ -53,14 +54,16 @@ def clean_environment(cls):
     for key in filter_keys:
       del os.environ[key]
 
-  def __init__(self, pex=sys.argv[0], interpreter=None, env=ENV):
+  def __init__(self, pex=sys.argv[0], interpreter=None, env=ENV, verify_entry_point=False):
     self._pex = pex
     self._interpreter = interpreter or PythonInterpreter.get()
     self._pex_info = PexInfo.from_pex(self._pex)
     self._pex_info_overrides = PexInfo.from_env(env=env)
     self._vars = env
     self._envs = []
     self._working_set = None
+    if verify_entry_point:
+      self._do_entry_point_verification()
 
   def _activate(self):
     if not self._working_set:
@@ -520,3 +523,32 @@ def run(self, args=(), with_chroot=False, blocking=True, setsid=False, **kwargs)
                                     stderr=kwargs.pop('stderr', None),
                                     **kwargs)
     return process.wait() if blocking else process
+
+  def _do_entry_point_verification(self):
+
+    entry_point = self._pex_info.entry_point
+    ep_split = entry_point.split(':')
+
+    # a.b.c:m ->
+    # ep_module = 'a.b.c'
+    # ep_method = 'm'
+
+    # Only module is specified
+    if len(ep_split) == 1:
+      ep_module = ep_split[0]
+      import_statement = 'import {}'.format(ep_module)
+    elif len(ep_split) == 2:
+      ep_module = ep_split[0]
+      ep_method = ep_split[1]
+      import_statement = 'from {} import {}'.format(ep_module, ep_method)
+    else:
+      raise self.InvalidEntryPoint("Failed to parse: `{}`".format(entry_point))
+
+    with named_temporary_file() as fp:
+      fp.write(import_statement.encode('utf-8'))
+      fp.close()
+      retcode = self.run([fp.name], env={'PEX_INTERPRETER': '1'})
+      if retcode != 0:
+        raise self.InvalidEntryPoint('Invalid entry point: `{}`\n'
+                                'Entry point verification failed: `{}`'
+                                .format(entry_point, import_statement))

tests/test_integration.py

@@ -959,3 +959,25 @@ def test_pex_resource_bundling():
 
       assert rc == 0
       assert stdout == b'hello\n'
+
+
+@pytest.mark.skipif(IS_PYPY)
+def test_entry_point_verification_3rdparty():
+  with temporary_dir() as td:
+    pex_out_path = os.path.join(td, 'pex.pex')
+    res = run_pex_command(['Pillow==5.2.0',
+                           '-e', 'PIL:Image',
+                           '-o', pex_out_path,
+                           '--validate-entry-point'])
+    res.assert_success()
+
+
+@pytest.mark.skipif(IS_PYPY)
+def test_invalid_entry_point_verification_3rdparty():
+  with temporary_dir() as td:
+    pex_out_path = os.path.join(td, 'pex.pex')
+    res = run_pex_command(['Pillow==5.2.0',
+                           '-e', 'PIL:invalid',
+                           '-o', pex_out_path,
+                           '--validate-entry-point'])
+    res.assert_failure()

tests/test_pex.py

@@ -4,13 +4,16 @@
 import os
 import sys
 import textwrap
+from contextlib import contextmanager
 from types import ModuleType
 
 import pytest
+from twitter.common.contextutil import temporary_file
 
 from pex.compatibility import WINDOWS, nested, to_bytes
 from pex.installer import EggInstaller, WheelInstaller
 from pex.pex import PEX
+from pex.pex_builder import PEXBuilder
 from pex.testing import (
     make_installer,
     named_temporary_file,
@@ -254,3 +257,55 @@ def test_pex_paths():
 
       fake_stdout.seek(0)
       assert fake_stdout.read() == b'42'
+
+
+@contextmanager
+def _add_test_hello_to_pex(ep):
+  with temporary_dir() as td:
+    hello_file = "\n".join([
+      "def hello():",
+      "  print('hello')",
+    ])
+    with temporary_file(root_dir=td) as tf:
+      with open(tf.name, 'w') as handle:
+        handle.write(hello_file)
+
+      pex_builder = PEXBuilder()
+      pex_builder.add_source(tf.name, 'test.py')
+      pex_builder.set_entry_point(ep)
+      pex_builder.freeze()
+      yield pex_builder
+
+
+def test_pex_verify_entry_point_method_should_pass():
+  with _add_test_hello_to_pex('test:hello') as pex_builder:
+    # No error should happen here because `test:hello` is correct
+    PEX(pex_builder.path(),
+        interpreter=pex_builder.interpreter,
+        verify_entry_point=True)
+
+
+def test_pex_verify_entry_point_module_should_pass():
+  with _add_test_hello_to_pex('test') as pex_builder:
+    # No error should happen here because `test` is correct
+    PEX(pex_builder.path(),
+        interpreter=pex_builder.interpreter,
+        verify_entry_point=True)
+
+
+def test_pex_verify_entry_point_method_should_fail():
+  with _add_test_hello_to_pex('test:invalid_entry_point') as pex_builder:
+    # Expect InvalidEntryPoint due to invalid entry point method
+    with pytest.raises(PEX.InvalidEntryPoint):
+      PEX(pex_builder.path(),
+          interpreter=pex_builder.interpreter,
+          verify_entry_point=True)
+
+
+def test_pex_verify_entry_point_module_should_fail():
+  with _add_test_hello_to_pex('invalid.module') as pex_builder:
+    # Expect InvalidEntryPoint due to invalid entry point module
+    with pytest.raises(PEX.InvalidEntryPoint):
+      PEX(pex_builder.path(),
+          interpreter=pex_builder.interpreter,
+          verify_entry_point=True)