Repository: pex-tool/pex

Issue:
The `sys.excepthook` can contain scrubbed code in a zipapp PEX leading to errors.
See an example over in https://github.com/pantsbuild/pants/issues/15877.

PEX should ensure that various `sys` facilities that have been modified with venv code (through `.pth` files, `site.py` and the like) are reset to their pristine states as part of `sys` scrubbing.


Discussion:

Resolving patch:
diff --git a/pex/pex.py b/pex/pex.py
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -408,6 +408,11 @@ def patch_all(path, path_importer_cache, modules):
             patch_dict(sys.path_importer_cache, path_importer_cache)
             patch_dict(sys.modules, modules)
 
+            # N.B.: These hooks may contain custom code set via site.py or `.pth` files that we've
+            # just scrubbed. Reset them to factory defaults to avoid scrubbed code.
+            sys.displayhook = sys.__displayhook__
+            sys.excepthook = sys.__excepthook__  # type: ignore[assignment] # Builtin typeshed bug.
+
         new_sys_path, new_sys_path_importer_cache, new_sys_modules = self.minimum_sys(inherit_path)
 
         if self._vars.PEX_EXTRA_SYS_PATH:


Test patch:
diff --git a/pex/testing.py b/pex/testing.py
--- a/pex/testing.py
+++ b/pex/testing.py
@@ -463,7 +463,7 @@ def bootstrap_python_installer(dest):
 
 
 def ensure_python_distribution(version):
-    # type: (str) -> Tuple[str, str, Callable[[Iterable[str]], Text]]
+    # type: (str) -> Tuple[str, str, str, Callable[[Iterable[str]], Text]]
     if version not in ALL_PY_VERSIONS:
         raise ValueError("Please constrain version to one of {}".format(ALL_PY_VERSIONS))
 
@@ -522,11 +522,11 @@ def run_pyenv(args):
         # type: (Iterable[str]) -> Text
         return to_unicode(subprocess.check_output([pyenv] + list(args), env=pyenv_env))
 
-    return python, pip, run_pyenv
+    return interpreter_location, python, pip, run_pyenv
 
 
 def ensure_python_venv(version, latest_pip=True, system_site_packages=False):
-    python, pip, _ = ensure_python_distribution(version)
+    _, python, pip, _ = ensure_python_distribution(version)
     venv = safe_mkdtemp()
     if version in _ALL_PY3_VERSIONS:
         args = [python, "-m", "venv", venv]
@@ -547,7 +547,7 @@ def ensure_python_venv(version, latest_pip=True, system_site_packages=False):
 
 def ensure_python_interpreter(version):
     # type: (str) -> str
-    python, _, _ = ensure_python_distribution(version)
+    _, python, _, _ = ensure_python_distribution(version)
     return python
 
 
diff --git a/tests/integration/test_issue_1809.py b/tests/integration/test_issue_1809.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/test_issue_1809.py
@@ -0,0 +1,91 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import os.path
+import shutil
+import subprocess
+from textwrap import dedent
+
+from pex.common import safe_open
+from pex.testing import PY310, ensure_python_distribution, make_project, run_pex_command
+from pex.typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+def test_excepthook_scrubbing(tmpdir):
+    # type: (Any) -> None
+
+    original_python_installation, original_python, _, _ = ensure_python_distribution(PY310)
+
+    custom_python_installation = os.path.join(str(tmpdir), "custom")
+    shutil.copytree(original_python_installation, custom_python_installation)
+    custom_python = os.path.join(custom_python_installation, "bin", "python")
+
+    project_dir = os.path.join(str(tmpdir), "custom_excepthook")
+    with safe_open(os.path.join(project_dir, "custom_excepthook.py"), "w") as fp:
+        fp.write("def custom_excepthook(typ, val, _tb): print('EXC: {} {}'.format(typ, val))")
+    with safe_open(os.path.join(project_dir, "custom_excepthook.pth"), "w") as fp:
+        fp.write("import custom_excepthook; sys.excepthook = custom_excepthook.custom_excepthook")
+
+    with safe_open(os.path.join(project_dir, "setup.py"), "w") as fp:
+        fp.write(
+            dedent(
+                """\
+                from distutils import sysconfig
+
+                from setuptools import setup
+
+
+                setup(
+                    name="custom_excepthook",
+                    version="0.0.1",
+                    py_modules=["custom_excepthook"],
+                    data_files=[(sysconfig.get_python_lib(), ["custom_excepthook.pth"])]
+                )
+                """
+            )
+        )
+
+    subprocess.check_call(args=[custom_python, "-m", "pip", "install", project_dir])
+
+    src = os.path.join(str(tmpdir), "src")
+    with safe_open(os.path.join(src, "app.py"), "w") as fp:
+        fp.write(
+            dedent(
+                """\
+                import sys
+
+
+                if getattr(sys.excepthook, "__name__", None) == "custom_excepthook":
+                    import custom_excepthook
+
+                print("SUCCESS")
+                """
+            )
+        )
+
+    pex = os.path.join(str(tmpdir), "pex")
+    create_pex_args = ["-D", src, "-m", "app", "-o", pex]
+
+    # Ensure this complicated test setup reproduces the original error in
+    # https://github.com/pantsbuild/pex/issues/1809
+    # (via https://github.com/pantsbuild/pants/issues/15877).
+    run_pex_command(args=["pex==2.1.92", "-c", "pex", "--"] + create_pex_args).assert_success()
+
+    # A Python installation without the custom excepthook installed via .pth should work just fine.
+    assert "SUCCESS\n" == subprocess.check_output(args=[original_python, pex]).decode("utf-8")
+
+    # But with the custom excepthook installed pre-PEX boot, we should reproduce earlier failures.
+    process = subprocess.Popen(args=[custom_python, pex], stdout=subprocess.PIPE)
+    stdout, _ = process.communicate()
+    assert 0 != process.returncode
+    assert (
+        "EXC: <class 'ModuleNotFoundError'> No module named 'custom_excepthook'\n"
+        == stdout.decode("utf-8")
+    )
+
+    # Now demonstrate the fix.
+    run_pex_command(args=create_pex_args).assert_success()
+    assert "SUCCESS\n" == subprocess.check_output(args=[custom_python, pex]).decode("utf-8")
diff --git a/tests/test_interpreter.py b/tests/test_interpreter.py
--- a/tests/test_interpreter.py
+++ b/tests/test_interpreter.py
@@ -218,7 +218,7 @@ def test_iter_interpreter_path_filter_symlink(self, test_interpreter1, test_inte
 
     def test_pyenv_shims(self, tmpdir):
         # type: (Any) -> None
-        py37, _, run_pyenv = ensure_python_distribution(PY37)
+        _, py37, _, run_pyenv = ensure_python_distribution(PY37)
         py310 = ensure_python_interpreter(PY310)
 
         pyenv_root = str(run_pyenv(["root"]).strip())