From ce7f8d85aa50d043282d53aec55c6b23dbc6bba3 Mon Sep 17 00:00:00 2001 From: John Sirois Date: Wed, 4 May 2022 16:57:13 -0700 Subject: [PATCH] Complete support for PEX Python interpreter mode. (#1748) Apply @BrandonTheBuilder's work in #1746 to the `--venv` case and flesh out complete support for Python interpreter emulation by adding `.` to the `sys.path` when running as an interpreter and handling interpreter execution of directories. Fixes #1745 --- .gitignore | 1 + pex/common.py | 2 +- pex/pex.py | 70 ++++++--- pex/venv/pex.py | 33 +++- tests/integration/venv_ITs/test_issue_1745.py | 145 ++++++++++++++++++ tests/test_pex.py | 6 +- 6 files changed, 227 insertions(+), 30 deletions(-) create mode 100644 tests/integration/venv_ITs/test_issue_1745.py diff --git a/.gitignore b/.gitignore index 5782df013..02a1912d0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ *~ *.pyc +*.pyo *.egg-info /.cache /venv diff --git a/pex/common.py b/pex/common.py index 548aab0ae..5a4344796 100644 --- a/pex/common.py +++ b/pex/common.py @@ -73,7 +73,7 @@ def filter_pyc_files(files): for f in files: # For Python 2.7, `.pyc` files are compiled as siblings to `.py` files (there is no # __pycache__ dir). - if not f.endswith(".pyc") and not is_pyc_temporary_file(f): + if not f.endswith((".pyc", ".pyo")) and not is_pyc_temporary_file(f): yield f diff --git a/pex/pex.py b/pex/pex.py index f0b1c73c5..9962adf98 100644 --- a/pex/pex.py +++ b/pex/pex.py @@ -32,7 +32,20 @@ from pex.variables import ENV, Variables if TYPE_CHECKING: - from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Set, Tuple, TypeVar + from typing import ( + Any, + Dict, + Iterable, + Iterator, + List, + Mapping, + NoReturn, + Optional, + Set, + Tuple, + TypeVar, + Union, + ) _K = TypeVar("_K") _V = TypeVar("_V") @@ -593,20 +606,23 @@ def log(msg, V=1): def execute_interpreter(self): # type: () -> Any + + # A Python interpreter always inserts the CWD at the head of the sys.path. + sys.path.insert(0, ".") + args = sys.argv[1:] - options = [] + python_options = [] for index, arg in enumerate(args): - # Check if the arg is an expected startup arg - if arg.startswith("-") and arg not in {"-", "-c", "-m"}: - options.append(arg) - continue + # Check if the arg is an expected startup arg. + if arg.startswith("-") and arg not in ("-", "-c", "-m"): + python_options.append(arg) else: args = args[index:] break - # The pex was called with interpreter options - if options: - return self.execute_with_options(options, args) + # The pex was called with Python interpreter options + if python_options: + return self.execute_with_options(python_options, args) if args: # NB: We take care here to setup sys.argv to match how CPython does it for each case. @@ -624,7 +640,8 @@ def execute_interpreter(self): if arg == "-": content = sys.stdin.read() else: - with open(arg) as fp: + file_path = arg if os.path.isfile(arg) else os.path.join(arg, "__main__.py") + with open(file_path) as fp: content = fp.read() except IOError as e: return "Could not open {} in the environment [{}]: {}".format( @@ -640,26 +657,29 @@ def execute_interpreter(self): code.interact() return None - def execute_with_options(self, options, args): + def execute_with_options( + self, + python_options, # type: List[str] + args, # List[str] + ): + # type: (...) -> Union[NoReturn, Any] """ Restart the process passing the given options to the python interpreter """ - # Find the pex bootstrap location - # We expect this environment variable to be set by __main__.py + # Find the installed (unzipped) PEX entry point. main = sys.modules.get("__main__") - if main: - run_pex = os.path.dirname(main.__file__) - interpreter = PythonInterpreter.get().binary - cmdline = [interpreter] + options + [run_pex] + args - TRACER.log( - "Re-executing with interpreter options: " - "cmdline={cmdline!r}, ".format( - cmdline=" ".join(cmdline), - ) + if not main or not main.__file__: + # N.B.: This should never happen. + return "Unable to resolve PEX __main__ module file: {}".format(main) + + python = sys.executable + cmdline = [python] + python_options + [main.__file__] + args + TRACER.log( + "Re-executing with Python interpreter options: cmdline={cmdline!r}".format( + cmdline=" ".join(cmdline) ) - os.execv(interpreter, cmdline) - else: - return "Unable to resolve pex path" + ) + os.execv(python, cmdline) def execute_script(self, script_name): # type: (str) -> Any diff --git a/pex/venv/pex.py b/pex/venv/pex.py index 2acf6b56e..3357a3c00 100644 --- a/pex/venv/pex.py +++ b/pex/venv/pex.py @@ -447,6 +447,36 @@ def iter_valid_venv_pythons(): ) if entry_point == PEX_INTERPRETER_ENTRYPOINT and len(sys.argv) > 1: args = sys.argv[1:] + + python_options = [] + for index, arg in enumerate(args): + # Check if the arg is an expected startup arg + if arg.startswith("-") and arg not in ("-", "-c", "-m"): + python_options.append(arg) + else: + args = args[index:] + break + + # The pex was called with Python interpreter options, so we need to re-exec to + # respect those: + if python_options: + # Find the installed (unzipped) PEX entry point. + main = sys.modules.get("__main__") + if not main or not main.__file__: + # N.B.: This should never happen. + sys.stderr.write( + "Unable to resolve PEX __main__ module file: {{}}\\n".format(main) + ) + sys.exit(1) + + python = sys.executable + cmdline = [python] + python_options + [main.__file__] + args + sys.stderr.write( + "Re-executing with Python interpreter options: " + "cmdline={{cmdline!r}}\\n".format(cmdline=" ".join(cmdline)) + ) + os.execv(python, cmdline) + arg = args[0] if arg == "-m": if len(args) < 2: @@ -468,7 +498,8 @@ def iter_valid_venv_pythons(): elif arg == "-": content = sys.stdin.read() else: - with open(arg) as fp: + file_path = arg if os.path.isfile(arg) else os.path.join(arg, "__main__.py") + with open(file_path) as fp: content = fp.read() ast = compile(content, filename, "exec", flags=0, dont_inherit=1) diff --git a/tests/integration/venv_ITs/test_issue_1745.py b/tests/integration/venv_ITs/test_issue_1745.py new file mode 100644 index 000000000..61e635dad --- /dev/null +++ b/tests/integration/venv_ITs/test_issue_1745.py @@ -0,0 +1,145 @@ +# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +import os +import subprocess +from textwrap import dedent + +import pytest +from colors import colors + +from pex.common import safe_open +from pex.enum import Enum +from pex.testing import IntegResults, run_pex_command +from pex.typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, List, Optional + + import attr # vendor:skip +else: + from pex.third_party import attr + + +# N.B.: To test that Python interpreter options are forwarded to Python we include an "assert False" +# in the code executed that we would normally expect to fail. Only by adding the `-O` option, which +# causes assertions to be ignored by Python, should the code execute fully without failing. +# See: https://docs.python.org/3/using/cmdline.html#cmdoption-O + +CODE = dedent( + """\ + import sys + + import colors + + + assert False, colors.red("Failed") + print(colors.green("Worked: {}".format(" ".join(sys.argv[1:])))) + """ +) + + +@attr.s(frozen=True) +class ExecutionConfiguration(object): + args = attr.ib(factory=list) # type: List[str] + cwd = attr.ib(default=None) # type: Optional[str] + stdin = attr.ib(default=None) # type: Optional[bytes] + + +class PythonInterfaceOption(Enum["PythonInterfaceOption.Value"]): + class Value(Enum.Value): + pass + + DASHC = Value("-c ") + DASHM = Value("-m ") + DASH = Value("- ( from STDIN)") + FILE = Value("") + DIRECTORY = Value("") + + +@pytest.fixture +def execution_configuration( + request, # type: Any + tmpdir, # type: Any +): + # type: (...) -> ExecutionConfiguration + if request.param is PythonInterfaceOption.DASHC: + return ExecutionConfiguration(args=["-c", CODE]) + + if request.param is PythonInterfaceOption.DASH: + return ExecutionConfiguration(args=["-"], stdin=CODE.encode("utf-8")) + + if request.param not in ( + PythonInterfaceOption.DASHM, + PythonInterfaceOption.FILE, + PythonInterfaceOption.DIRECTORY, + ): + raise AssertionError("Unexpected PythonInterfaceOption: {}".format(request.param)) + + src = os.path.join(str(tmpdir), "src") + python_file = os.path.join( + src, + "{filename}.py".format( + filename="__main__" if request.param is PythonInterfaceOption.DIRECTORY else "module" + ), + ) + with safe_open(python_file, "w") as fp: + fp.write(CODE) + + if request.param is PythonInterfaceOption.DASHM: + return ExecutionConfiguration(args=["-m", "module"], cwd=src) + + if request.param is PythonInterfaceOption.FILE: + return ExecutionConfiguration(args=[python_file]) + + return ExecutionConfiguration(args=[src]) + + +@pytest.mark.parametrize( + "execution_mode_args", [pytest.param([], id="ZIPAPP"), pytest.param(["--venv"], id="VENV")] +) +@pytest.mark.parametrize( + "execution_configuration", + [pytest.param(value, id=str(value)) for value in PythonInterfaceOption.values()], + indirect=True, +) +def test_interpreter_mode_python_options( + execution_mode_args, # type: List[str] + execution_configuration, # type: ExecutionConfiguration + tmpdir, # type: Any +): + # type: (...) -> None + + pex = os.path.join(str(tmpdir), "pex") + run_pex_command(args=["ansicolors==1.1.8", "-o", pex] + execution_mode_args).assert_success() + + def execute_pex(disable_assertions): + # type: (bool) -> IntegResults + args = [pex] + if disable_assertions: + args.append("-O") + args.extend(execution_configuration.args) + args.extend(("program", "args")) + process = subprocess.Popen( + args=args, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + cwd=execution_configuration.cwd, + ) + stdout, stderr = process.communicate(input=execution_configuration.stdin) + return IntegResults( + output=stdout.decode("utf-8"), + error=stderr.decode("utf-8"), + return_code=process.returncode, + ) + + # With no `-O` for Python, the CODE assertion should fail. + result = execute_pex(disable_assertions=False) + result.assert_failure() + assert colors.red("Failed") in result.error + + # But with `-O` forwarded to Python, the CODE assertion should be skipped. + result = execute_pex(disable_assertions=True) + result.assert_success() + assert colors.green("Worked: program args") in result.output diff --git a/tests/test_pex.py b/tests/test_pex.py index 4d1fc49c6..b0a912b38 100644 --- a/tests/test_pex.py +++ b/tests/test_pex.py @@ -607,7 +607,7 @@ def test_execute_interpreter_dashc_program(): assert b"" == stderr -def test_execute_interpreter_dashc_program_with_option(): +def test_execute_interpreter_dashc_program_with_python_options(): with temporary_dir() as pex_chroot: pex_builder = PEXBuilder(path=pex_chroot) pex_builder.freeze() @@ -663,7 +663,7 @@ def test_execute_interpreter_dashm_module(): assert b"" == stderr -def test_execute_interpreter_dashm_module_with_option(): +def test_execute_interpreter_dashm_module_with_python_options(): # type: () -> None with temporary_dir() as pex_chroot: pex_builder = PEXBuilder(path=pex_chroot) @@ -722,7 +722,7 @@ def test_execute_interpreter_stdin_program(): assert b"" == stderr -def test_execute_interpreter_stdin_program_with_option(): +def test_execute_interpreter_stdin_program_with_python_options(): # type: () -> None with temporary_dir() as pex_chroot: pex_builder = PEXBuilder(path=pex_chroot)