Set PEX_ROOT for export post-processing command. (cherry-pick of #17177…

…) (#17200)

This is both a leak that foils our recommendations to users for maintenance of 
Pants caches and results in unneeded extra work when an export finds an 
empty ~/.pex but would have had hits in ~/.cache/pants/named_caches/pex_root 
if it had been pointed there.

A reapplication of part of #17177, but not really a straight cherry-pick, as it would not 
apply even remotely cleanly.

src/python/pants/backend/python/goals/export.py

@@ -14,7 +14,7 @@
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.pex import Pex, PexRequest
 from pants.backend.python.util_rules.pex_cli import PexPEX
-from pants.backend.python.util_rules.pex_environment import PythonExecutable
+from pants.backend.python.util_rules.pex_environment import PexEnvironment, PythonExecutable
 from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest
 from pants.core.goals.export import (
     ExportError,
@@ -79,7 +79,7 @@ def debug_hint(self) -> str | None:
 
 @rule
 async def export_virtualenv(
-    request: _ExportVenvRequest, python_setup: PythonSetup, pex_pex: PexPEX
+    request: _ExportVenvRequest, python_setup: PythonSetup, pex_pex: PexPEX, pex_env: PexEnvironment
 ) -> ExportResult:
     if request.resolve:
         interpreter_constraints = InterpreterConstraints(
@@ -127,6 +127,7 @@ async def export_virtualenv(
     merged_digest = await Get(Digest, MergeDigests([pex_pex.digest, requirements_pex.digest]))
     pex_pex_path = os.path.join("{digest_root}", pex_pex.exe)
     maybe_resolve_str = f"for the resolve '{request.resolve}' " if request.resolve else ""
+    complete_pex_env = pex_env.in_workspace()
     return ExportResult(
         f"virtualenv {maybe_resolve_str}(using Python {py_version})",
         dest,
@@ -143,15 +144,20 @@ async def export_virtualenv(
                     "--remove=all",
                     f"{{digest_root}}/{py_version}",
                 ],
-                {"PEX_MODULE": "pex.tools"},
+                {
+                    **complete_pex_env.environment_dict(python_configured=True),
+                    "PEX_MODULE": "pex.tools",
+                },
             ),
             PostProcessingCommand(["rm", "-f", pex_pex_path]),
         ],
     )
 
 
 @rule
-async def export_tool(request: ExportPythonTool, pex_pex: PexPEX) -> ExportResult:
+async def export_tool(
+    request: ExportPythonTool, pex_pex: PexPEX, pex_env: PexEnvironment
+) -> ExportResult:
     assert request.pex_request is not None
 
     # TODO: Unify export_virtualenv() and export_tool(), since their implementations mostly overlap.
@@ -173,6 +179,7 @@ async def export_tool(request: ExportPythonTool, pex_pex: PexPEX) -> ExportResul
     pex_pex_digest = await Get(Digest, AddPrefix(pex_pex.digest, pex_pex_dir))
 
     merged_digest = await Get(Digest, MergeDigests([pex_pex_digest, pex.digest]))
+    complete_pex_env = pex_env.in_workspace()
     return ExportResult(
         f"virtualenv for the tool '{request.resolve_name}'",
         dest,
@@ -188,7 +195,10 @@ async def export_tool(request: ExportPythonTool, pex_pex: PexPEX) -> ExportResul
                     "--remove=all",
                     f"{{digest_root}}/{request.resolve_name}",
                 ],
-                {"PEX_MODULE": "pex.tools"},
+                {
+                    **complete_pex_env.environment_dict(python_configured=True),
+                    "PEX_MODULE": "pex.tools",
+                },
             ),
             PostProcessingCommand(["rm", "-rf", pex_pex_dest]),
         ],

src/python/pants/backend/python/goals/export_test.py

@@ -87,7 +87,8 @@ def run(enable_resolves: bool) -> ExportResults:
                 "--remove=all",
                 f"{{digest_root}}/{current_interpreter}",
             )
-            assert ppc0.extra_env == FrozenDict({"PEX_MODULE": "pex.tools"})
+            assert ppc0.extra_env["PEX_MODULE"] == "pex.tools"
+            assert "PEX_ROOT" in ppc0.extra_env
 
             ppc1 = result.post_processing_cmds[1]
             assert ppc1.argv == (