Add helper to extract distribution info from a VenvPex (#11665)

This is useful for querying what version of a dist is being used, such as is required by #11658.

[ci skip-rust]
[ci skip-build-wheels]

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

@@ -6,6 +6,7 @@
 import dataclasses
 import functools
 import itertools
+import json
 import logging
 import shlex
 from collections import defaultdict
@@ -14,6 +15,8 @@
 from textwrap import dedent
 from typing import FrozenSet, Iterable, List, Mapping, Sequence, Set, Tuple, TypeVar
 
+import packaging.specifiers
+import packaging.version
 from pkg_resources import Requirement
 from typing_extensions import Protocol
 
@@ -28,7 +31,7 @@
     PythonExecutable,
 )
 from pants.engine.addresses import Address
-from pants.engine.collection import DeduplicatedCollection
+from pants.engine.collection import Collection, DeduplicatedCollection
 from pants.engine.engine_aware import EngineAwareParameter
 from pants.engine.fs import (
     EMPTY_DIGEST,
@@ -738,6 +741,7 @@ def python(self, bash: BashBinary, pex_environment: PexEnvironment) -> VenvScrip
 @dataclass(frozen=True)
 class VenvPex:
     digest: Digest
+    pex_filename: str
     pex: Script
     python: Script
     bin: FrozenDict[str, Script]
@@ -825,6 +829,7 @@ async def create_venv_pex(
 
     return VenvPex(
         digest=input_digest,
+        pex_filename=result.pex_filename,
         pex=pex.script,
         python=python.script,
         bin=FrozenDict((bin_name, venv_script.script) for bin_name, venv_script in scripts.items()),
@@ -1010,5 +1015,50 @@ async def setup_venv_pex_process(request: VenvPexProcess) -> Process:
     )
 
 
+@dataclass(frozen=True)
+class PexDistributionInfo:
+    """Information about an individual distribution in a PEX file, as reported by `PEX_TOOLS=1
+    repository info -v`."""
+
+    project_name: str
+    version: packaging.version.Version
+    requires_python: packaging.specifiers.SpecifierSet | None
+    requires_dists: tuple[Requirement, ...]
+
+
+class PexResolveInfo(Collection[PexDistributionInfo]):
+    """Information about all distributions resolved in a PEX file, as reported by `PEX_TOOLS=1
+    repository info -v`."""
+
+
+@rule
+async def determine_venv_pex_resolve_info(venv_pex: VenvPex) -> PexResolveInfo:
+    process_result = await Get(
+        ProcessResult,
+        VenvPexProcess(
+            venv_pex,
+            argv=["repository", "info", "-v"],
+            extra_env={"PEX_TOOLS": "1"},
+            input_digest=venv_pex.digest,
+            description=f"Determine distributions found in {venv_pex.pex_filename}",
+            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))
+
+
 def rules():
     return [*collect_rules(), *pex_cli.rules()]

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

@@ -11,6 +11,8 @@
 from typing import Dict, Iterable, Iterator, List, Mapping, Tuple, cast
 
 import pytest
+from packaging.specifiers import SpecifierSet
+from packaging.version import Version
 from pkg_resources import Requirement
 
 from pants.backend.python.target_types import (
@@ -20,11 +22,13 @@
 )
 from pants.backend.python.util_rules.pex import (
     Pex,
+    PexDistributionInfo,
     PexInterpreterConstraints,
     PexPlatforms,
     PexProcess,
     PexRequest,
     PexRequirements,
+    PexResolveInfo,
     VenvPex,
     VenvPexProcess,
 )
@@ -321,6 +325,7 @@ def rule_runner() -> RuleRunner:
             QueryRule(Process, (PexProcess,)),
             QueryRule(Process, (VenvPexProcess,)),
             QueryRule(ProcessResult, (Process,)),
+            QueryRule(PexResolveInfo, (VenvPex,)),
         ]
     )
 
@@ -607,3 +612,19 @@ def test_additional_inputs(rule_runner: RuleRunner) -> None:
         with zipfp.open("__main__.py", "r") as main:
             main_content = main.read().decode()
     assert main_content[: len(preamble)] == preamble
+
+
+def test_venv_pex_resolve_info(rule_runner: RuleRunner) -> None:
+    venv_pex = create_pex_and_get_all_data(
+        rule_runner, pex_type=VenvPex, 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[2] == PexDistributionInfo(
+        "idna", Version("2.10"), SpecifierSet("!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,>=2.7"), ()
+    )
+    assert dists[3].project_name == "requests"
+    assert dists[3].version == Version("2.23.0")
+    assert Requirement.parse('PySocks!=1.5.7,>=1.5.6; extra == "socks"') in dists[3].requires_dists
+    assert dists[4].project_name == "urllib3"