Repository: pex-tool/pex

Issue:
Support exporting PYTHONPATH before running user code
As discussed in Slack, it'd be nice to use `PYTHONPATH` in Pants (and not have it be scrubbed) when invoking user PEX script so that the user code can use `subprocess.run([sys.executable, ...], ...)` and not have to pass a constructed `PYTHONPATH` in the `env=` argument.

How this looks is dealer's choice, as it's hidden from me, the client, through Pants.


Discussion:
It's really a bug that the `PEX_EXTRA_SYS_PATH` additions are not inherited by sub-processes. This should probably occur automatically since its how a normal Python interpreter works and PEX files strive to act like normal Python interpreters with the deps baked in.
When the target of a `subprocess.run` is itself a PEX, there exists `--no-strip-pex-env` to compose with this.
The history here is #989

Resolving patch:
diff --git a/pex/pex.py b/pex/pex.py
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -413,7 +413,19 @@ def patch_all(path, path_importer_cache, modules):
 
         if self._vars.PEX_EXTRA_SYS_PATH:
             TRACER.log("Adding %s to sys.path" % self._vars.PEX_EXTRA_SYS_PATH)
-            new_sys_path.extend(self._vars.PEX_EXTRA_SYS_PATH.split(":"))
+            extra_sys_path = self._vars.PEX_EXTRA_SYS_PATH.split(":")
+            new_sys_path.extend(extra_sys_path)
+
+            # Let Python subprocesses see the same sys.path additions we see. If those Python
+            # subprocesses are PEX subprocesses, they'll do their own (re-)scrubbing as needed.
+            if inherit_path is InheritPath.FALSE:
+                pythonpath_entries = extra_sys_path
+            else:
+                raw_pythonpath = os.environ.get(self._PYTHONPATH)
+                pythonpath = raw_pythonpath.split(os.pathsep) if raw_pythonpath else []
+                pythonpath_entries = pythonpath + extra_sys_path
+            os.environ[self._PYTHONPATH] = os.pathsep.join(pythonpath_entries)
+
         TRACER.log("New sys.path: %s" % new_sys_path)
 
         patch_all(new_sys_path, new_sys_path_importer_cache, new_sys_modules)
@@ -580,7 +592,7 @@ def execute_interpreter(self):
         # type: () -> Any
 
         # A Python interpreter always inserts the CWD at the head of the sys.path.
-        sys.path.insert(0, ".")
+        sys.path.insert(0, "")
 
         args = sys.argv[1:]
         python_options = []
diff --git a/pex/venv/pex.py b/pex/venv/pex.py
--- a/pex/venv/pex.py
+++ b/pex/venv/pex.py
@@ -417,7 +417,17 @@ def sys_executable_paths():
                 )
 
             os.environ["VIRTUAL_ENV"] = venv_dir
-            sys.path.extend(os.environ.get("PEX_EXTRA_SYS_PATH", "").split(os.pathsep))
+
+            # A Python interpreter always inserts the CWD at the head of the sys.path.
+            sys.path.insert(0, "")
+
+            pex_extra_sys_path = os.environ.get("PEX_EXTRA_SYS_PATH")
+            if pex_extra_sys_path:
+                sys.path.extend(pex_extra_sys_path.split(":"))
+
+                # Let Python subprocesses see the same sys.path additions we see. If those Python
+                # subprocesses are PEX subprocesses, they'll do their own (re-)scrubbing as needed.
+                os.environ["PYTHONPATH"] = pex_extra_sys_path
 
             bin_path = os.environ.get("PEX_VENV_BIN_PATH", {bin_path!r})
             if bin_path != "false":


Test patch:
diff --git a/tests/integration/test_issue_1825.py b/tests/integration/test_issue_1825.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/test_issue_1825.py
@@ -0,0 +1,157 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import json
+import os.path
+import subprocess
+from textwrap import dedent
+
+import pytest
+
+from pex.common import safe_open
+from pex.inherit_path import InheritPath
+from pex.orderedset import OrderedSet
+from pex.testing import make_env, run_pex_command
+from pex.typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from typing import Any, List
+
+
+DUMP_SYS_PATH_CODE = dedent(
+    """\
+    import json
+    import sys
+
+
+    json.dump(sys.path, sys.stdout)
+    """
+)
+
+
+def create_sys_path_dump_pex(
+    tmpdir,  # type: Any
+    *additional_args  # type: str
+):
+    # type: (...) -> str
+
+    exe = os.path.join(str(tmpdir), "exe.py")
+    with open(exe, "w") as fp:
+        fp.write(DUMP_SYS_PATH_CODE)
+
+    pex = os.path.join(str(tmpdir), "pex")
+    run_pex_command(args=["--exe", exe, "-o", pex] + list(additional_args)).assert_success()
+    return pex
+
+
+def execute_sys_path_dump_pex(
+    pex,  # type: str
+    *additional_args,  # type: str
+    **additional_env  # type: Any
+):
+    # type: (...) -> OrderedSet[str]
+
+    return OrderedSet(
+        os.path.realpath(entry)
+        for entry in cast(
+            "List[str]",
+            json.loads(
+                subprocess.check_output(
+                    args=[pex] + list(additional_args), env=make_env(**additional_env)
+                )
+            ),
+        )
+    )
+
+
+def read_additional_sys_path(
+    pex,  # type: str
+    *additional_args,  # type: str
+    **additional_env  # type: Any
+):
+    # type: (...) -> List[str]
+
+    isolated_sys_path = execute_sys_path_dump_pex(pex)
+    return list(
+        execute_sys_path_dump_pex(pex, *additional_args, **additional_env) - isolated_sys_path
+    )
+
+
+def assert_inherited(
+    tmpdir,  # type: Any
+    pex,  # type: str
+    *expected_inherited_sys_path_entries,  # type: str
+    **additional_env  # type: Any
+):
+    other_exe = os.path.join(str(tmpdir), "other_exe.py")
+    with safe_open(other_exe, "w") as fp:
+        fp.write(DUMP_SYS_PATH_CODE)
+
+    subprocess_proof = os.path.join(str(tmpdir), "proof")
+    assert not os.path.exists(subprocess_proof)
+
+    code = dedent(
+        """\
+        import subprocess
+        import sys
+
+
+        open({subprocess_proof!r}, "w").close()
+        sys.exit(subprocess.call([sys.executable, {other_exe!r}]))
+        """
+    ).format(subprocess_proof=subprocess_proof, other_exe=other_exe)
+    actual = read_additional_sys_path(pex, "-c", code, PEX_INTERPRETER=1, **additional_env)
+    assert os.path.exists(subprocess_proof)
+
+    previous = -1
+    for expected_inherited_sys_path_entry in expected_inherited_sys_path_entries:
+        assert expected_inherited_sys_path_entry in actual
+        current_index = actual.index(expected_inherited_sys_path_entry)
+        assert current_index > previous
+        previous = current_index
+
+
+@pytest.mark.parametrize(
+    "execution_mode_args",
+    [
+        pytest.param([], id="ZIPAPP"),
+        pytest.param(["--venv"], id="VENV"),
+    ],
+)
+def test_pex_run_extra_sys_path(
+    tmpdir,  # type: Any
+    execution_mode_args,  # type: List[str]
+):
+    # type: (...) -> None
+
+    pex = create_sys_path_dump_pex(tmpdir, *execution_mode_args)
+
+    foo = os.path.join(str(tmpdir), "foo")
+    assert [foo] == read_additional_sys_path(pex, PEX_EXTRA_SYS_PATH=foo)
+
+    bar = os.path.join(str(tmpdir), "bar")
+    assert_inherited(tmpdir, pex, bar, PEX_EXTRA_SYS_PATH=bar)
+
+
+@pytest.mark.parametrize(
+    "inherit_path",
+    [pytest.param(inherit_path, id=str(inherit_path)) for inherit_path in InheritPath.values()],
+)
+def test_pex_run_inherit_path_and_extra_sys_path(
+    tmpdir,  # type: Any
+    inherit_path,  # type: InheritPath.Value
+):
+    # type: (...) -> None
+
+    pex = create_sys_path_dump_pex(
+        tmpdir, "--inherit-path={inherit_path}".format(inherit_path=inherit_path)
+    )
+
+    esp = os.path.join(str(tmpdir), "esp")
+    pp1 = os.path.join(str(tmpdir), "pp1")
+    pp2 = os.path.join(str(tmpdir), "pp2")
+
+    expected = [esp] if inherit_path is InheritPath.FALSE else [pp1, pp2, esp]
+    assert_inherited(
+        tmpdir, pex, *expected, PEX_EXTRA_SYS_PATH=esp, PYTHONPATH=":".join((pp1, pp2))
+    )