Refactor venv export

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

@@ -12,9 +12,8 @@
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import PythonResolveField
 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 import Pex, PexProcess, 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_from_targets import RequirementsPexRequest
 from pants.core.goals.export import (
     ExportError,
@@ -28,8 +27,8 @@
 from pants.engine.environment import EnvironmentName
 from pants.engine.internals.native_engine import AddPrefix, Digest, MergeDigests
 from pants.engine.internals.selectors import Get, MultiGet
-from pants.engine.process import Process, ProcessResult
-from pants.engine.rules import collect_rules, rule
+from pants.engine.process import ProcessResult
+from pants.engine.rules import collect_rules, rule, rule_helper
 from pants.engine.target import Target
 from pants.engine.unions import UnionMembership, UnionRule, union
 from pants.util.docutil import bin_name
@@ -78,122 +77,130 @@ def debug_hint(self) -> str | None:
         return self.resolve_name
 
 
-@rule
-async def export_virtualenv(
-    request: _ExportVenvRequest, python_setup: PythonSetup, pex_pex: PexPEX
+@rule_helper
+async def _do_export(
+    requirements_pex: Pex,
+    pex_pex: PexPEX,
+    dest: str,
+    resolve_name: str,
+    qualify_path_with_python_version: bool,
 ) -> ExportResult:
-    if request.resolve:
-        interpreter_constraints = InterpreterConstraints(
-            python_setup.resolves_to_interpreter_constraints.get(
-                request.resolve, python_setup.interpreter_constraints
-            )
-        )
-    else:
-        interpreter_constraints = InterpreterConstraints.create_from_targets(
-            request.root_python_targets, python_setup
-        ) or InterpreterConstraints(python_setup.interpreter_constraints)
-
-    requirements_pex = await Get(
-        Pex,
-        RequirementsPexRequest(
-            (tgt.address for tgt in request.root_python_targets),
-            hardcoded_interpreter_constraints=interpreter_constraints,
-        ),
-    )
-
-    # Note that an internal-only pex will always have the `python` field set.
-    # See the build_pex() rule in pex.py.
-    interpreter = cast(PythonExecutable, requirements_pex.python)
+    if requirements_pex.python is None:
+        # An internal-only pex will always have the `python` field set.
+        # See the build_pex() rule and _determine_pex_python_and_platforms() helper in pex.py.
+        raise ExportError(f"The PEX to be exported for {resolve_name} must be internal_only.")
 
-    # Get the full python version (including patch #), so we can use it as the venv name.
+    # Get the full python version (including patch #).
     res = await Get(
         ProcessResult,
-        Process(
+        PexProcess(
+            pex=requirements_pex,
             description="Get interpreter version",
             argv=[
-                interpreter.path,
                 "-c",
                 "import sys; print('.'.join(str(x) for x in sys.version_info[0:3]))",
             ],
+            extra_env={"PEX_INTERPRETER": "1"},
         ),
     )
     py_version = res.stdout.strip().decode()
 
-    dest = (
-        os.path.join("python", "virtualenvs", path_safe(request.resolve))
-        if request.resolve
-        else os.path.join("python", "virtualenv")
-    )
+    # NOTE: We add a unique prefix to the pex_pex path to avoid conflicts when multiple
+    # venvs are concurrently exporting. Without this prefix all the invocations write
+    # the pex_pex to `python/virtualenvs/tools/pex`, and the `rm -f` of the pex_pex
+    # path in one export will delete the binary out from under the others.
+    pex_pex_dir = f".{resolve_name}.tmp"
+    pex_pex_digest = await Get(Digest, AddPrefix(pex_pex.digest, pex_pex_dir))
+    pex_pex_dest = os.path.join("{digest_root}", pex_pex_dir)
+
+    merged_digest = await Get(Digest, MergeDigests([pex_pex_digest, requirements_pex.digest]))
 
-    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 ""
+    description = f"for {resolve_name} " if resolve_name else ""
     return ExportResult(
-        f"virtualenv {maybe_resolve_str}(using Python {py_version})",
+        f"virtualenv {description}(using Python {py_version})",
         dest,
         digest=merged_digest,
         post_processing_cmds=[
             PostProcessingCommand(
                 [
-                    interpreter.path,
-                    pex_pex_path,
+                    requirements_pex.python.path,
+                    os.path.join(pex_pex_dest, pex_pex.exe),
                     os.path.join("{digest_root}", requirements_pex.name),
                     "venv",
                     "--pip",
                     "--collisions-ok",
                     "--remove=all",
-                    f"{{digest_root}}/{py_version}",
+                    f"{{digest_root}}/{py_version if qualify_path_with_python_version else ''}",
                 ],
                 {"PEX_MODULE": "pex.tools"},
             ),
-            PostProcessingCommand(["rm", "-f", pex_pex_path]),
+            # Remove the PEX pex, to avoid confusion.
+            PostProcessingCommand(["rm", "-rf", pex_pex_dest]),
         ],
     )
 
 
+@rule
+async def export_virtualenv_for_targets(
+    request: _ExportVenvRequest,
+    python_setup: PythonSetup,
+    pex_pex: PexPEX,
+) -> ExportResult:
+    if request.resolve:
+        interpreter_constraints = InterpreterConstraints(
+            python_setup.resolves_to_interpreter_constraints.get(
+                request.resolve, python_setup.interpreter_constraints
+            )
+        )
+    else:
+        interpreter_constraints = InterpreterConstraints.create_from_targets(
+            request.root_python_targets, python_setup
+        ) or InterpreterConstraints(python_setup.interpreter_constraints)
+
+    # Note that a pex created from a RequirementsPexRequest has packed layout, which should lead
+    # to the best performance in this use case.
+    requirements_pex = await Get(
+        Pex,
+        RequirementsPexRequest(
+            (tgt.address for tgt in request.root_python_targets),
+            hardcoded_interpreter_constraints=interpreter_constraints,
+        ),
+    )
+
+    dest = (
+        os.path.join("python", "virtualenvs", path_safe(request.resolve))
+        if request.resolve
+        else os.path.join("python", "virtualenv")
+    )
+
+    export_result = await _do_export(
+        requirements_pex,
+        pex_pex,
+        dest,
+        request.resolve or "",
+        qualify_path_with_python_version=True,
+    )
+    return export_result
+
+
 @rule
 async def export_tool(request: ExportPythonTool, pex_pex: PexPEX) -> ExportResult:
     assert request.pex_request is not None
 
-    # TODO: Unify export_virtualenv() and export_tool(), since their implementations mostly overlap.
-    dest = os.path.join("python", "virtualenvs", "tools")
+    # TODO: It seems unnecessary to qualify with "tools", since the tool resolve names don't collide
+    #  with user resolve names.  We should get rid of this via a deprecation cycle.
+    dest = os.path.join("python", "virtualenvs", "tools", request.resolve_name)
     pex = await Get(Pex, PexRequest, request.pex_request)
-    if not request.pex_request.internal_only:
-        raise ExportError(f"The PexRequest for {request.resolve_name} must be internal_only.")
-
-    # Note that an internal-only pex will always have the `python` field set.
-    # See the build_pex() rule in pex.py.
-    interpreter = cast(PythonExecutable, pex.python)
-
-    # NOTE: We add a unique-per-tool prefix to the pex_pex path to avoid conflicts when
-    # multiple tools are concurrently exporting. Without this prefix all the `export_tool`
-    # invocations write the pex_pex to `python/virtualenvs/tools/pex`, and the `rm -f` of
-    # the pex_pex path in one export will delete the binary out from under the others.
-    pex_pex_dir = f".{request.resolve_name}.tmp"
-    pex_pex_dest = os.path.join("{digest_root}", pex_pex_dir)
-    pex_pex_digest = await Get(Digest, AddPrefix(pex_pex.digest, pex_pex_dir))
-
-    merged_digest = await Get(Digest, MergeDigests([pex_pex_digest, pex.digest]))
-    return ExportResult(
-        f"virtualenv for the tool '{request.resolve_name}'",
+    export_result = await _do_export(
+        pex,
+        pex_pex,
         dest,
-        digest=merged_digest,
-        post_processing_cmds=[
-            PostProcessingCommand(
-                [
-                    interpreter.path,
-                    os.path.join(pex_pex_dest, pex_pex.exe),
-                    os.path.join("{digest_root}", pex.name),
-                    "venv",
-                    "--collisions-ok",
-                    "--remove=all",
-                    f"{{digest_root}}/{request.resolve_name}",
-                ],
-                {"PEX_MODULE": "pex.tools"},
-            ),
-            PostProcessingCommand(["rm", "-rf", pex_pex_dest]),
-        ],
+        request.resolve_name,
+        # TODO: It is pretty ad-hoc that we do add the interpreter version for resolves but not for tools.
+        #  We should pick one and deprecate the other.
+        qualify_path_with_python_version=False,
     )
+    return export_result
 
 
 @rule

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

@@ -29,6 +29,7 @@ class _ToolConfig:
     version: str
     experimental: bool = False
     backend_prefix: str | None = "lint"
+    takes_ics: bool = True
 
     @property
     def package(self) -> str:
@@ -38,16 +39,16 @@ def package(self) -> str:
 EXPORTED_TOOLS: List[_ToolConfig] = [
     _ToolConfig(name="add-trailing-comma", version="2.2.3", experimental=True),
     _ToolConfig(name="autoflake", version="1.3.1", experimental=True),
-    _ToolConfig(name="bandit", version="1.6.2"),
+    _ToolConfig(name="bandit", version="1.6.2", takes_ics=False),
     _ToolConfig(name="black", version="22.3.0"),
     _ToolConfig(name="docformatter", version="1.3.1"),
-    _ToolConfig(name="flake8", version="4.0.1"),
+    _ToolConfig(name="flake8", version="4.0.1", takes_ics=False),
     _ToolConfig(name="isort", version="5.10.1"),
-    _ToolConfig(name="pylint", version="2.13.1"),
+    _ToolConfig(name="pylint", version="2.13.1", takes_ics=False),
     _ToolConfig(name="pyupgrade", version="2.31.1", experimental=True),
     _ToolConfig(name="yapf", version="0.32.0"),
     _ToolConfig(name="mypy", version="0.940", backend_prefix="typecheck"),
-    _ToolConfig(name="pytest", version="7.1.0", backend_prefix=None),
+    _ToolConfig(name="pytest", version="7.1.0", backend_prefix=None, takes_ics=False),
 ]
 
 
@@ -70,6 +71,8 @@ def build_config(tmpdir: str) -> Mapping:
             "version": f"{tool_config.name}=={tool_config.version}",
             "lockfile": f"{tmpdir}/3rdparty/{tool_config.name}.lock",
         }
+        if tool_config.takes_ics:
+            cfg[tool_config.name]["interpreter_constraints"] = (f"=={platform.python_version()}",)
 
         if not tool_config.backend_prefix:
             continue
@@ -110,6 +113,7 @@ def test_export() -> None:
 
         # NOTE: Not every tool implements --version so this is the best we can do.
         lib_dir = os.path.join(export_dir, "lib", f"python{py_minor_version}", "site-packages")
+
         expected_tool_dir = os.path.join(
             lib_dir, f"{tool_config.package}-{tool_config.version}.dist-info"
         )

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

@@ -72,14 +72,14 @@ def run(enable_resolves: bool) -> ExportResults:
         )
         all_results = rule_runner.request(ExportResults, [ExportVenvsRequest(targets)])
 
-        for result in all_results:
+        for result, resolve in zip(all_results, ["a", "b"] if enable_resolves else [""]):
             assert len(result.post_processing_cmds) == 2
 
             ppc0 = result.post_processing_cmds[0]
             assert ppc0.argv[1:] == (
                 # The first arg is the full path to the python interpreter, which we
                 # don't easily know here, so we ignore it in this comparison.
-                os.path.join("{digest_root}", ".", "pex"),
+                os.path.join("{digest_root}", f".{resolve}.tmp", ".", "pex"),
                 os.path.join("{digest_root}", "requirements.pex"),
                 "venv",
                 "--pip",
@@ -92,8 +92,8 @@ def run(enable_resolves: bool) -> ExportResults:
             ppc1 = result.post_processing_cmds[1]
             assert ppc1.argv == (
                 "rm",
-                "-f",
-                os.path.join("{digest_root}", ".", "pex"),
+                "-rf",
+                os.path.join("{digest_root}", f".{resolve}.tmp"),
             )
             assert ppc1.extra_env == FrozenDict()
 

src/python/pants/backend/python/util_rules/pex.py

@@ -55,7 +55,6 @@
 from pants.engine.fs import EMPTY_DIGEST, AddPrefix, CreateDigest, Digest, FileContent, MergeDigests
 from pants.engine.internals.native_engine import Snapshot
 from pants.engine.internals.selectors import MultiGet
-from pants.engine.platform import Platform
 from pants.engine.process import Process, ProcessCacheScope, ProcessResult
 from pants.engine.rules import Get, collect_rules, rule, rule_helper
 from pants.engine.target import HydratedSources, HydrateSourcesRequest, SourcesField, Targets
@@ -502,7 +501,7 @@ async def _setup_pex_requirements(
 
 @rule(level=LogLevel.DEBUG)
 async def build_pex(
-    request: PexRequest, python_setup: PythonSetup, platform: Platform, pex_subsystem: PexSubsystem
+    request: PexRequest, python_setup: PythonSetup, pex_subsystem: PexSubsystem
 ) -> BuildPexResult:
     """Returns a PEX with the given settings."""
     argv = [

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

@@ -43,7 +43,7 @@ class ExportRequest:
 @frozen_after_init
 @dataclass(unsafe_hash=True)
 class PostProcessingCommand:
-    """A command to run as a local processe after an exported digest is materialized."""
+    """A command to run as a local process after an exported digest is materialized."""
 
     # Values in the argv tuple can contain the format specifier "{digest_root}", which will be
     # substituted with the (absolute) path to the location under distdir in which the
@@ -71,11 +71,6 @@ class ExportResult:
     # Materialize this digest.
     digest: Digest
     # Run these commands as local processes after the digest is materialized.
-    # Values in each args string tuple can contain the format specifier "{digest_root}", which
-    # will be substituted with the (absolute) path to the location under distdir in which the
-    # digest is materialized.
-    # Each command will be run with an environment consistent of just PATH, set to the Pants
-    # process's own PATH env var.
     post_processing_cmds: tuple[PostProcessingCommand, ...]
 
     def __init__(