Support --venv mode for internal PEXes.

src/python/pants/backend/awslambda/python/rules.py

@@ -12,13 +12,13 @@
 )
 from pants.backend.python.util_rules import pex_from_targets
 from pants.backend.python.util_rules.pex import (
-    Pex,
     PexInterpreterConstraints,
     PexPlatforms,
-    PexProcess,
     PexRequest,
     PexRequirements,
     TwoStepPex,
+    VenvPex,
+    VenvPexProcess,
 )
 from pants.backend.python.util_rules.pex_from_targets import (
     PexFromTargetsRequest,
@@ -30,7 +30,6 @@
     OutputPathField,
     PackageFieldSet,
 )
-from pants.engine.fs import Digest, MergeDigests
 from pants.engine.process import ProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.unions import UnionRule
@@ -92,19 +91,18 @@ async def package_python_awslambda(
     )
 
     lambdex_pex, pex_result, handler = await MultiGet(
-        Get(Pex, PexRequest, lambdex_request),
+        Get(VenvPex, PexRequest, lambdex_request),
         Get(TwoStepPex, TwoStepPexFromTargetsRequest, pex_request),
         Get(ResolvedPythonAwsHandler, ResolvePythonAwsHandlerRequest(field_set.handler)),
     )
-    input_digest = await Get(Digest, MergeDigests((pex_result.pex.digest, lambdex_pex.digest)))
 
     # NB: Lambdex modifies its input pex in-place, so the input file is also the output file.
     result = await Get(
         ProcessResult,
-        PexProcess(
+        VenvPexProcess(
             lambdex_pex,
             argv=("build", "-e", handler.val, output_filename),
-            input_digest=input_digest,
+            input_digest=pex_result.pex.digest,
             output_files=(output_filename,),
             description=f"Setting up handler in {output_filename}",
         ),

src/python/pants/backend/codegen/protobuf/python/rules.py

@@ -9,13 +9,13 @@
 from pants.backend.codegen.protobuf.python.python_protobuf_subsystem import PythonProtobufSubsystem
 from pants.backend.codegen.protobuf.target_types import ProtobufGrpcToggle, ProtobufSources
 from pants.backend.python.target_types import PythonSources
-from pants.backend.python.util_rules import extract_pex, pex
-from pants.backend.python.util_rules.extract_pex import ExtractedPexDistributions
+from pants.backend.python.util_rules import pex
 from pants.backend.python.util_rules.pex import (
-    Pex,
     PexInterpreterConstraints,
     PexRequest,
     PexRequirements,
+    VenvPex,
+    VenvPexRequest,
 )
 from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
 from pants.core.util_rules.source_files import SourceFilesRequest
@@ -95,23 +95,27 @@ async def generate_python_from_protobuf(
         target_stripped_sources_request,
     )
 
-    # To run the MyPy Protobuf plugin, we first install it with Pex, then extract the wheels and
-    # point Protoc to the extracted wheels with its `--plugin` argument.
-    extracted_mypy_wheels = None
+    protoc_gen_mypy_script = "protoc-gen-mypy"
+    mypy_pex = None
     if python_protobuf_subsystem.mypy_plugin:
         mypy_pex = await Get(
-            Pex,
-            PexRequest(
-                output_filename="mypy_protobuf.pex",
-                internal_only=True,
-                requirements=PexRequirements([python_protobuf_subsystem.mypy_plugin_version]),
-                # This is solely to ensure that we use an appropriate interpreter when resolving
-                # the distribution. We don't actually run the distribution directly with Python,
-                # as we extract out its binary.
-                interpreter_constraints=PexInterpreterConstraints(["CPython>=3.5"]),
+            VenvPex,
+            VenvPexRequest(
+                bin_names=[protoc_gen_mypy_script],
+                pex_request=PexRequest(
+                    output_filename="mypy_protobuf.pex",
+                    internal_only=True,
+                    requirements=PexRequirements([python_protobuf_subsystem.mypy_plugin_version]),
+                    # TODO(John Sirois): Fix these interpreter constraints to track the actual
+                    #  python requirement of the mypy_plugin_version or else plumb an option for
+                    #  manually setting the constraint to track what mypy_plugin_version needs:
+                    #  https://github.com/pantsbuild/pants/issues/11565
+                    # Here we guess a constraint that will likely work with any mypy_plugin_version
+                    # selected.
+                    interpreter_constraints=PexInterpreterConstraints(["CPython>=3.5"]),
+                ),
             ),
         )
-        extracted_mypy_wheels = await Get(ExtractedPexDistributions, Pex, mypy_pex)
 
     downloaded_grpc_plugin = (
         await Get(
@@ -128,22 +132,17 @@ async def generate_python_from_protobuf(
         downloaded_protoc_binary.digest,
         empty_output_dir,
     ]
-    if extracted_mypy_wheels:
-        unmerged_digests.append(extracted_mypy_wheels.digest)
+    if mypy_pex:
+        unmerged_digests.append(mypy_pex.digest)
     if downloaded_grpc_plugin:
         unmerged_digests.append(downloaded_grpc_plugin.digest)
     input_digest = await Get(Digest, MergeDigests(unmerged_digests))
 
     argv = [downloaded_protoc_binary.exe, "--python_out", output_dir]
-    if extracted_mypy_wheels:
-        mypy_plugin_path = next(
-            p
-            for p in extracted_mypy_wheels.wheel_directory_paths
-            if p.startswith(".deps/mypy_protobuf-")
-        )
+    if mypy_pex:
         argv.extend(
             [
-                f"--plugin=protoc-gen-mypy={mypy_plugin_path}/bin/protoc-gen-mypy",
+                f"--plugin=protoc-gen-mypy={mypy_pex.bin[protoc_gen_mypy_script].argv0}",
                 "--mypy_out",
                 output_dir,
             ]
@@ -154,15 +153,10 @@ async def generate_python_from_protobuf(
         )
     argv.extend(target_sources_stripped.snapshot.files)
 
-    env = {}
-    if extracted_mypy_wheels:
-        env["PYTHONPATH"] = ":".join(extracted_mypy_wheels.wheel_directory_paths)
-
     result = await Get(
         ProcessResult,
         Process(
             argv,
-            env=env,
             input_digest=input_digest,
             description=f"Generating Python sources from {request.protocol_target.address}.",
             level=LogLevel.DEBUG,
@@ -196,7 +190,6 @@ async def generate_python_from_protobuf(
 def rules():
     return [
         *collect_rules(),
-        *extract_pex.rules(),
         *pex.rules(),
         UnionRule(GenerateSourcesRequest, GeneratePythonFromProtobufRequest),
     ]

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

@@ -12,11 +12,11 @@
 
 from pants.backend.python.subsystems.python_tool_base import PythonToolBase
 from pants.backend.python.util_rules.pex import (
-    Pex,
     PexInterpreterConstraints,
-    PexProcess,
     PexRequest,
     PexRequirements,
+    VenvPex,
+    VenvPexProcess,
 )
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFiles,
@@ -216,13 +216,13 @@ async def create_coverage_config(coverage: CoverageSubsystem) -> CoverageConfig:
 
 @dataclass(frozen=True)
 class CoverageSetup:
-    pex: Pex
+    pex: VenvPex
 
 
 @rule
 async def setup_coverage(coverage: CoverageSubsystem) -> CoverageSetup:
     pex = await Get(
-        Pex,
+        VenvPex,
         PexRequest(
             output_filename="coverage.pex",
             internal_only=True,
@@ -250,11 +250,11 @@ async def merge_coverage_data(
         Get(Digest, AddPrefix(data.digest, prefix=data.address.path_safe_spec))
         for data in data_collection
     )
-    input_digest = await Get(Digest, MergeDigests((*coverage_digests, coverage_setup.pex.digest)))
+    input_digest = await Get(Digest, MergeDigests(coverage_digests))
     prefixes = sorted(f"{data.address.path_safe_spec}/.coverage" for data in data_collection)
     result = await Get(
         ProcessResult,
-        PexProcess(
+        VenvPexProcess(
             coverage_setup.pex,
             argv=("combine", *prefixes),
             input_digest=input_digest,
@@ -291,7 +291,6 @@ async def generate_coverage_reports(
             (
                 merged_coverage_data.coverage_data,
                 coverage_config.digest,
-                coverage_setup.pex.digest,
                 sources.source_files.snapshot.digest,
             )
         ),
@@ -319,7 +318,7 @@ async def generate_coverage_reports(
             else None
         )
         pex_processes.append(
-            PexProcess(
+            VenvPexProcess(
                 coverage_setup.pex,
                 argv=(report_type.report_name,),
                 input_digest=input_digest,
@@ -329,7 +328,9 @@ async def generate_coverage_reports(
                 level=LogLevel.DEBUG,
             )
         )
-    results = await MultiGet(Get(ProcessResult, PexProcess, process) for process in pex_processes)
+    results = await MultiGet(
+        Get(ProcessResult, VenvPexProcess, process) for process in pex_processes
+    )
     result_stdouts = tuple(res.stdout for res in results)
     result_snapshots = await MultiGet(Get(Snapshot, Digest, res.output_digest) for res in results)
 

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

@@ -21,9 +21,10 @@
 from pants.backend.python.util_rules.pex import (
     Pex,
     PexInterpreterConstraints,
-    PexProcess,
     PexRequest,
     PexRequirements,
+    VenvPex,
+    VenvPexProcess,
 )
 from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
 from pants.backend.python.util_rules.python_sources import (
@@ -115,7 +116,6 @@ async def setup_pytest_for_target(
         all_targets, python_setup
     )
 
-    # Defaults to zip_safe=False.
     requirements_pex_request = Get(
         Pex,
         PexFromTargetsRequest,
@@ -128,23 +128,7 @@ async def setup_pytest_for_target(
             output_filename="pytest.pex",
             requirements=PexRequirements(pytest.get_requirement_strings()),
             interpreter_constraints=interpreter_constraints,
-            entry_point="pytest:main",
             internal_only=True,
-            additional_args=(
-                # NB: We set `--not-zip-safe` because Pytest plugin discovery, which uses
-                # `importlib_metadata` and thus `zipp`, does not play nicely when doing import
-                # magic directly from zip files. `zipp` has pathologically bad behavior with large
-                # zipfiles.
-                # TODO: this does have a performance cost as the pex must now be expanded to disk.
-                # Long term, it would be better to fix Zipp (whose fix would then need to be used
-                # by importlib_metadata and then by Pytest). See
-                # https://github.com/jaraco/zipp/pull/26.
-                "--not-zip-safe",
-                # TODO(John Sirois): Support shading python binaries:
-                #   https://github.com/pantsbuild/pants/issues/9206
-                "--pex-path",
-                requirements_pex_request.input.output_filename,
-            ),
         ),
     )
 
@@ -183,14 +167,23 @@ async def setup_pytest_for_target(
         field_set_source_files_request,
     )
 
+    pytest_runner_pex = await Get(
+        VenvPex,
+        PexRequest(
+            output_filename="pytest_runner.pex",
+            interpreter_constraints=interpreter_constraints,
+            entry_point="pytest:main",
+            internal_only=True,
+            pex_path=[pytest_pex, requirements_pex],
+        ),
+    )
+
     input_digest = await Get(
         Digest,
         MergeDigests(
             (
                 coverage_config.digest,
                 prepared_sources.source_files.snapshot.digest,
-                requirements_pex.digest,
-                pytest_pex.digest,
                 *(binary.digest for binary in assets),
             )
         ),
@@ -227,8 +220,8 @@ async def setup_pytest_for_target(
     cache_scope = ProcessCacheScope.NEVER if test_subsystem.force else ProcessCacheScope.SUCCESSFUL
     process = await Get(
         Process,
-        PexProcess(
-            pytest_pex,
+        VenvPexProcess(
+            pytest_runner_pex,
             argv=(*pytest.options.args, *coverage_args, *field_set_source_files.files),
             extra_env=extra_env,
             input_digest=input_digest,

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

@@ -24,11 +24,11 @@
     SetupPyCommandsField,
 )
 from pants.backend.python.util_rules.pex import (
-    Pex,
     PexInterpreterConstraints,
-    PexProcess,
     PexRequest,
     PexRequirements,
+    VenvPex,
+    VenvPexProcess,
 )
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFilesRequest,
@@ -411,7 +411,7 @@ async def run_setup_py(req: RunSetupPyRequest, setuptools: Setuptools) -> RunSet
     # Note that this pex has no entrypoint. We use it to run our generated setup.py, which
     # in turn imports from and invokes setuptools.
     setuptools_pex = await Get(
-        Pex,
+        VenvPex,
         PexRequest(
             output_filename="setuptools.pex",
             internal_only=True,
@@ -423,16 +423,15 @@ async def run_setup_py(req: RunSetupPyRequest, setuptools: Setuptools) -> RunSet
             ),
         ),
     )
-    input_digest = await Get(Digest, MergeDigests((req.chroot.digest, setuptools_pex.digest)))
     # The setuptools dist dir, created by it under the chroot (not to be confused with
     # pants's own dist dir, at the buildroot).
     dist_dir = "dist/"
     result = await Get(
         ProcessResult,
-        PexProcess(
+        VenvPexProcess(
             setuptools_pex,
             argv=("setup.py", *req.args),
-            input_digest=input_digest,
+            input_digest=req.chroot.digest,
             # setuptools commands that create dists write them to the distdir.
             # TODO: Could there be other useful files to capture?
             output_directories=(dist_dir,),

src/python/pants/backend/python/lint/bandit/rules.py

@@ -8,11 +8,11 @@
 from pants.backend.python.target_types import InterpreterConstraintsField, PythonSources
 from pants.backend.python.util_rules import pex
 from pants.backend.python.util_rules.pex import (
-    Pex,
     PexInterpreterConstraints,
-    PexProcess,
     PexRequest,
     PexRequirements,
+    VenvPex,
+    VenvPexProcess,
 )
 from pants.core.goals.lint import LintReport, LintRequest, LintResult, LintResults, LintSubsystem
 from pants.core.util_rules import stripped_source_files
@@ -63,7 +63,7 @@ async def bandit_lint_partition(
     partition: BanditPartition, bandit: Bandit, lint_subsystem: LintSubsystem
 ) -> LintResult:
     bandit_pex_request = Get(
-        Pex,
+        VenvPex,
         PexRequest(
             output_filename="bandit.pex",
             internal_only=True,
@@ -90,15 +90,13 @@ async def bandit_lint_partition(
         bandit_pex_request, config_digest_request, source_files_request
     )
 
-    input_digest = await Get(
-        Digest, MergeDigests((source_files.snapshot.digest, bandit_pex.digest, config_digest))
-    )
+    input_digest = await Get(Digest, MergeDigests((source_files.snapshot.digest, config_digest)))
 
     report_file_name = "bandit_report.txt" if lint_subsystem.reports_dir else None
 
     result = await Get(
         FallibleProcessResult,
-        PexProcess(
+        VenvPexProcess(
             bandit_pex,
             argv=generate_args(
                 source_files=source_files, bandit=bandit, report_file_name=report_file_name