Support extracting PexResolveInfo from a Pex.

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

@@ -13,7 +13,7 @@
 from dataclasses import dataclass
 from pathlib import PurePath
 from textwrap import dedent
-from typing import FrozenSet, Iterable, List, Mapping, Sequence, Set, Tuple, TypeVar
+from typing import FrozenSet, Iterable, Iterator, List, Mapping, Sequence, Set, Tuple, TypeVar
 
 import packaging.specifiers
 import packaging.version
@@ -24,7 +24,7 @@
 from pants.backend.python.target_types import PexPlatformsField as PythonPlatformsField
 from pants.backend.python.target_types import PythonRequirementsField
 from pants.backend.python.util_rules import pex_cli
-from pants.backend.python.util_rules.pex_cli import PexCliProcess
+from pants.backend.python.util_rules.pex_cli import PexCliProcess, PexPEX
 from pants.backend.python.util_rules.pex_environment import (
     PexEnvironment,
     PexRuntimeEnvironment,
@@ -1031,6 +1031,27 @@ class PexResolveInfo(Collection[PexDistributionInfo]):
     repository info -v`."""
 
 
+def parse_repository_info(repository_info: str) -> PexResolveInfo:
+    def iter_dist_info() -> Iterator[PexDistributionInfo]:
+        for line in repository_info.splitlines():
+            info = json.loads(line)
+            requires_python = info["requires_python"]
+            yield PexDistributionInfo(
+                project_name=info["project_name"],
+                version=packaging.version.Version(info["version"]),
+                requires_python=(
+                    packaging.specifiers.SpecifierSet(requires_python)
+                    if requires_python is not None
+                    else None
+                ),
+                requires_dists=tuple(
+                    Requirement.parse(req) for req in sorted(info["requires_dists"])
+                ),
+            )
+
+    return PexResolveInfo(sorted(iter_dist_info(), key=lambda dist: dist.project_name))
+
+
 @rule
 async def determine_venv_pex_resolve_info(venv_pex: VenvPex) -> PexResolveInfo:
     process_result = await Get(
@@ -1044,20 +1065,23 @@ async def determine_venv_pex_resolve_info(venv_pex: VenvPex) -> PexResolveInfo:
             level=LogLevel.DEBUG,
         ),
     )
-    dists = []
-    for line in process_result.stdout.decode().splitlines():
-        info = json.loads(line)
-        dists.append(
-            PexDistributionInfo(
-                project_name=info["project_name"],
-                version=packaging.version.Version(info["version"]),
-                requires_python=packaging.specifiers.SpecifierSet(info["requires_python"] or ""),
-                requires_dists=tuple(
-                    Requirement.parse(req) for req in sorted(info["requires_dists"])
-                ),
-            )
-        )
-    return PexResolveInfo(sorted(dists, key=lambda dist: dist.project_name))
+    return parse_repository_info(process_result.stdout.decode())
+
+
+@rule
+async def determine_pex_resolve_info(pex_pex: PexPEX, pex: Pex) -> PexResolveInfo:
+    process_result = await Get(
+        ProcessResult,
+        PexProcess(
+            pex=Pex(digest=pex_pex.digest, name=pex_pex.exe, python=pex.python),
+            argv=[pex.name, "repository", "info", "-v"],
+            input_digest=pex.digest,
+            extra_env={"PEX_MODULE": "pex.tools"},
+            description=f"Determine distributions found in {pex.name}",
+            level=LogLevel.DEBUG,
+        ),
+    )
+    return parse_repository_info(process_result.stdout.decode())
 
 
 def rules():

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

@@ -97,20 +97,29 @@ def __post_init__(self) -> None:
             raise ValueError("`--pex-root` flag not allowed. We set its value for you.")
 
 
+class PexPEX(DownloadedExternalTool):
+    """The Pex PEX binary."""
+
+
+@rule
+async def download_pex_pex(pex_binary: PexBinary) -> PexPEX:
+    pex_pex = await Get(
+        DownloadedExternalTool, ExternalToolRequest, pex_binary.get_request(Platform.current)
+    )
+    return PexPEX(digest=pex_pex.digest, exe=pex_pex.exe)
+
+
 @rule
 async def setup_pex_cli_process(
     request: PexCliProcess,
-    pex_binary: PexBinary,
+    pex_binary: PexPEX,
     pex_env: PexEnvironment,
     python_native_code: PythonNativeCode,
     global_options: GlobalOptions,
     pex_runtime_env: PexRuntimeEnvironment,
 ) -> Process:
     tmpdir = ".tmp"
-    gets: List[Get] = [
-        Get(DownloadedExternalTool, ExternalToolRequest, pex_binary.get_request(Platform.current)),
-        Get(Digest, CreateDigest([Directory(tmpdir)])),
-    ]
+    gets: List[Get] = [Get(Digest, CreateDigest([Directory(tmpdir)]))]
     cert_args = []
 
     # The certs file will typically not be in the repo, so we can't digest it via a PathGlobs.
@@ -127,15 +136,15 @@ async def setup_pex_cli_process(
         )
         cert_args = ["--cert", chrooted_ca_certs_path]
 
-    downloaded_pex_bin, *digests_to_merge = await MultiGet(gets)
-    digests_to_merge.append(downloaded_pex_bin.digest)
+    digests_to_merge = [pex_binary.digest]
+    digests_to_merge.extend(await MultiGet(gets))
     if request.additional_input_digest:
         digests_to_merge.append(request.additional_input_digest)
     input_digest = await Get(Digest, MergeDigests(digests_to_merge))
 
     pex_root_path = ".cache/pex_root"
     argv = [
-        downloaded_pex_bin.exe,
+        pex_binary.exe,
         *cert_args,
         "--python-path",
         create_path_env_var(pex_env.interpreter_search_paths),

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

@@ -33,6 +33,7 @@
     VenvPexProcess,
 )
 from pants.backend.python.util_rules.pex import rules as pex_rules
+from pants.backend.python.util_rules.pex_cli import PexPEX
 from pants.engine.addresses import Address
 from pants.engine.fs import CreateDigest, Digest, FileContent
 from pants.engine.process import Process, ProcessResult
@@ -325,7 +326,9 @@ def rule_runner() -> RuleRunner:
             QueryRule(Process, (PexProcess,)),
             QueryRule(Process, (VenvPexProcess,)),
             QueryRule(ProcessResult, (Process,)),
+            QueryRule(PexResolveInfo, (Pex,)),
             QueryRule(PexResolveInfo, (VenvPex,)),
+            QueryRule(PexPEX, ()),
         ]
     )
 
@@ -354,7 +357,7 @@ def create_pex_and_get_all_data(
         main=main,
         sources=sources,
         additional_inputs=additional_inputs,
-        additional_args=("--include-tools", *additional_pex_args),
+        additional_args=additional_pex_args,
     )
     rule_runner.set_options(
         ["--backend-packages=pants.backend.python", *additional_pants_args],
@@ -364,26 +367,37 @@ def create_pex_and_get_all_data(
     pex = rule_runner.request(pex_type, [request])
     if isinstance(pex, Pex):
         digest = pex.digest
+        pex_pex = rule_runner.request(PexPEX, [])
+        process = rule_runner.request(
+            Process,
+            [
+                PexProcess(
+                    Pex(digest=pex_pex.digest, name=pex_pex.exe, python=pex.python),
+                    argv=["-m", "pex.tools", pex.name, "info"],
+                    input_digest=pex.digest,
+                    extra_env=dict(PEX_INTERPRETER="1"),
+                    description="Extract PEX-INFO.",
+                )
+            ],
+        )
     elif isinstance(pex, VenvPex):
         digest = pex.digest
+        process = rule_runner.request(
+            Process,
+            [
+                VenvPexProcess(
+                    pex,
+                    argv=["info"],
+                    extra_env=dict(PEX_TOOLS="1"),
+                    description="Extract PEX-INFO.",
+                ),
+            ],
+        )
     else:
         raise AssertionError(f"Expected a Pex or a VenvPex but got a {type(pex)}.")
+
     rule_runner.scheduler.write_digest(digest)
     pex_path = os.path.join(rule_runner.build_root, "test.pex")
-
-    pex_process_type = PexProcess if isinstance(pex, Pex) else VenvPexProcess
-    process = rule_runner.request(
-        Process,
-        [
-            pex_process_type(
-                pex,
-                argv=["info"],
-                extra_env=dict(PEX_TOOLS="1"),
-                description="Extract PEX-INFO.",
-            ),
-        ],
-    )
-
     result = rule_runner.request(ProcessResult, [process])
     pex_info_content = result.stdout.decode()
 
@@ -614,13 +628,14 @@ def test_additional_inputs(rule_runner: RuleRunner) -> None:
     assert main_content[: len(preamble)] == preamble
 
 
-def test_venv_pex_resolve_info(rule_runner: RuleRunner) -> None:
+@pytest.mark.parametrize("pex_type", [Pex, VenvPex])
+def test_venv_pex_resolve_info(rule_runner: RuleRunner, pex_type: type[Pex | VenvPex]) -> None:
     venv_pex = create_pex_and_get_all_data(
-        rule_runner, pex_type=VenvPex, requirements=PexRequirements(["requests==2.23.0"])
+        rule_runner, pex_type=pex_type, requirements=PexRequirements(["requests==2.23.0"])
     )["pex"]
     dists = rule_runner.request(PexResolveInfo, [venv_pex])
-    assert dists[0] == PexDistributionInfo("certifi", Version("2020.12.5"), SpecifierSet(""), ())
-    assert dists[1] == PexDistributionInfo("chardet", Version("3.0.4"), SpecifierSet(""), ())
+    assert dists[0] == PexDistributionInfo("certifi", Version("2020.12.5"), None, ())
+    assert dists[1] == PexDistributionInfo("chardet", Version("3.0.4"), None, ())
     assert dists[2] == PexDistributionInfo(
         "idna", Version("2.10"), SpecifierSet("!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,>=2.7"), ()
     )