diff --git a/pex/bin/pex.py b/pex/bin/pex.py index dc99b70d0..8a6b2155c 100755 --- a/pex/bin/pex.py +++ b/pex/bin/pex.py @@ -355,6 +355,14 @@ def configure_clp(): help='The name of the generated .pex file: Omiting this will run PEX ' 'immediately and not save it to a file.') + parser.add_option( + '-p', '--preamble-file', + dest='preamble_file', + metavar='FILE', + default=None, + type=str, + help='The name of a file to be included as the preamble for the generated .pex file') + parser.add_option( '-r', '--requirement', dest='requirement_files', @@ -513,8 +521,15 @@ def build_pex(args, options, resolver_option_builder): if not interpreters: die('Could not find compatible interpreter', CANNOT_SETUP_INTERPRETER) + try: + with open(options.preamble_file) as preamble_fd: + preamble = preamble_fd.read() + except TypeError: + # options.preamble_file is None + preamble = None + interpreter = _lowest_version_interpreter(interpreters) - pex_builder = PEXBuilder(path=safe_mkdtemp(), interpreter=interpreter) + pex_builder = PEXBuilder(path=safe_mkdtemp(), interpreter=interpreter, preamble=preamble) pex_info = pex_builder.info pex_info.zip_safe = options.zip_safe diff --git a/tests/test_pex_binary.py b/tests/test_pex_binary.py index ccc6ccf70..ede3de1ec 100644 --- a/tests/test_pex_binary.py +++ b/tests/test_pex_binary.py @@ -3,8 +3,10 @@ from contextlib import contextmanager from optparse import OptionParser +from tempfile import NamedTemporaryFile -from pex.bin.pex import configure_clp, configure_clp_pex_resolution +from pex.bin.pex import build_pex, configure_clp, configure_clp_pex_resolution +from pex.compatibility import to_bytes from pex.fetcher import PyPIFetcher from pex.package import SourcePackage, WheelPackage from pex.resolver_options import ResolverOptionsBuilder @@ -88,6 +90,19 @@ def test_clp_constraints_txt(): assert options.constraint_files == ['requirements1.txt'] +def test_clp_preamble_file(): + with NamedTemporaryFile() as tmpfile: + tmpfile.write(to_bytes('print "foo!"')) + tmpfile.flush() + + parser, resolver_options_builder = configure_clp() + options, reqs = parser.parse_args(args=['--preamble-file', tmpfile.name]) + assert options.preamble_file == tmpfile.name + + pex_builder = build_pex(reqs, options, resolver_options_builder) + assert pex_builder._preamble == to_bytes('print "foo!"') + + def test_clp_prereleases(): with parser_pair() as (builder, parser): configure_clp_pex_resolution(parser, builder) diff --git a/tests/test_pex_builder.py b/tests/test_pex_builder.py index 66f022689..9521ab371 100644 --- a/tests/test_pex_builder.py +++ b/tests/test_pex_builder.py @@ -101,6 +101,30 @@ def builder(shebang): assert fp.read(len(expected_preamble)) == expected_preamble +def test_pex_builder_preamble(): + with temporary_dir() as td: + target = os.path.join(td, 'foo.pex') + should_create = os.path.join(td, 'foo.1') + + tempfile_preamble = "\n".join([ + "import sys", + "open('{0}', 'w').close()".format(should_create), + "sys.exit(3)" + ]) + + pex_builder = PEXBuilder(preamble=tempfile_preamble) + pex_builder.build(target) + + assert not os.path.exists(should_create) + + pex = PEX(target) + process = pex.run(blocking=False) + process.wait() + + assert process.returncode == 3 + assert os.path.exists(should_create) + + def test_pex_builder_compilation(): with nested(temporary_dir(), temporary_dir(), temporary_dir()) as (td1, td2, td3): src = os.path.join(td1, 'src.py')