diff --git a/pex/bin/pex.py b/pex/bin/pex.py index 8e0ce5804..0ef47b1f2 100755 --- a/pex/bin/pex.py +++ b/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__': diff --git a/pex/pex.py b/pex/pex.py index 112fa5ce9..02f2d642c 100644 --- a/pex/pex.py +++ b/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,7 +54,7 @@ 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) @@ -61,6 +62,8 @@ def __init__(self, pex=sys.argv[0], interpreter=None, 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)) diff --git a/tests/test_integration.py b/tests/test_integration.py index 1f53775f1..f8ac671a8 100644 --- a/tests/test_integration.py +++ b/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() diff --git a/tests/test_pex.py b/tests/test_pex.py index 22f2ef4a4..ac3d2017b 100644 --- a/tests/test_pex.py +++ b/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)