Relocate environment variable injection to before the interpreter is …

…run (#2260) As per #2224 --------- Co-authored-by: Erik Williamson <[email protected]>
pex-tool · Nov 7, 2023 · 4089022 · 4089022
1 parent 32e9116
commit 4089022
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 11 deletions.
diff --git a/pex/pex.py b/pex/pex.py
@@ -568,6 +568,9 @@ def _execute(self):
 
         self._clean_environment(strip_pex_env=self._pex_info.strip_pex_env)
 
+        for name, value in self._pex_info.inject_env.items():
+            os.environ.setdefault(name, value)
+
         if force_interpreter:
             TRACER.log("PEX_INTERPRETER specified, dropping into interpreter")
             return self.execute_interpreter()
@@ -596,8 +599,6 @@ def _execute(self):
                 EntryPoint.parse("run = {}".format(self._pex_info_overrides.entry_point))
             )
 
-        for name, value in self._pex_info.inject_env.items():
-            os.environ.setdefault(name, value)
         sys.argv[1:1] = list(self._pex_info.inject_args)
 
         if self._pex_info.script:

diff --git a/pex/venv/installer.py b/pex/venv/installer.py
@@ -816,7 +816,10 @@ def sys_executable_paths():
                                 "PEX_EXTRA_SYS_PATH"
                             )
                         del os.environ[key]
-
+  
+            for name, value in {inject_env!r}:
+                os.environ.setdefault(name, value)
+                
             pex_script = pex_overrides.get("PEX_SCRIPT") if pex_overrides else {script!r}
             if pex_script:
                 script_path = os.path.join(venv_bin_dir, pex_script)
@@ -829,7 +832,7 @@ def sys_executable_paths():
                 if pex_interpreter
                 else pex_overrides.get("PEX_MODULE", {entry_point!r} or PEX_INTERPRETER_ENTRYPOINT)
             )
-
+                                  
             if entry_point == PEX_INTERPRETER_ENTRYPOINT:
                 # A Python interpreter always inserts the CWD at the head of the sys.path.
                 # See https://docs.python.org/3/library/sys.html#sys.path
@@ -917,8 +920,6 @@ def sys_executable_paths():
                     sys.exit(0)
 
             if not is_exec_override:
-                for name, value in {inject_env!r}:
-                    os.environ.setdefault(name, value)
                 sys.argv[1:1] = {inject_args!r}
 
             module_name, _, function = entry_point.partition(":")

diff --git a/tests/integration/test_inject_env_and_args.py b/tests/integration/test_inject_env_and_args.py
@@ -72,11 +72,11 @@ def assert_FOO(
     assert_FOO(expected_env_value="baz", runtime_env_value="baz")
     assert_FOO(expected_env_value="", runtime_env_value="")
 
-    # Switching away from the built-in entrypoint should disable the injected env.
+    # Switching away from the built-in entrypoint should retain the injected env.
     assert (
-        "<not set>"
+        "bar"
         == subprocess.check_output(
-            args=[pex, "-c", print_FOO_env_code], env=make_env(PEX_INTERPRETER=1)
+            args=[pex, "-c", print_FOO_env_code], env=make_env(PEX_INTERPRETER=1, FOO="bar")
         )
         .decode("utf-8")
         .strip()
@@ -240,8 +240,8 @@ def assert_message(
     assert_message(b"Hello, world!")
     assert_message(b"42", MESSAGE="42")
 
-    # Switching away from the built-in entrypoint should disable injected args and env.
-    assert {"args": ["foo", "bar"], "MESSAGE": None} == json.loads(
+    # Switching away from the built-in entrypoint should disable injected args but not the env.
+    assert {"args": ["foo", "bar"], "MESSAGE": "Hello, world!"} == json.loads(
         subprocess.check_output(args=[pex, "foo", "bar"], env=make_env(PEX_MODULE="example"))
     )