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,5 +1,6 @@
 *~
 *.pyc
+*.pyo
 *.egg-info
 /.cache
 /venv

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
 
 

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

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)

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 <code>")
+    DASHM = Value("-m <module>")
+    DASH = Value("- (<code> from STDIN)")
+    FILE = Value("<python file>")
+    DIRECTORY = Value("<dir>")
+
+
+@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

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)