[internal] Add internal ability to generate lockfiles with PEX

src/python/pants/backend/java/lint/google_java_format/BUILD

@@ -3,6 +3,6 @@
 
 python_sources(dependencies=[":lockfile"])
 
-python_tests(name="tests")
+python_tests(name="tests", timeout=120)
 
 resource(name="lockfile", source="google_java_format.default.lockfile.txt")

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

@@ -26,6 +26,7 @@
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.lockfile_metadata import PythonLockfileMetadata
 from pants.backend.python.util_rules.pex import PexRequest, PexRequirements, VenvPex, VenvPexProcess
+from pants.backend.python.util_rules.pex_cli import PexCliProcess
 from pants.core.goals.generate_lockfiles import (
     GenerateLockfile,
     GenerateLockfileResult,
@@ -52,6 +53,7 @@
 class GeneratePythonLockfile(GenerateLockfile):
     requirements: FrozenOrderedSet[str]
     interpreter_constraints: InterpreterConstraints
+    use_pex: bool = False
     # Only kept for `[python].experimental_lockfile`, which is not using the new
     # "named resolve" semantics yet.
     _description: str | None = None
@@ -114,8 +116,8 @@ def warn_python_repos(option: str) -> None:
             "If lockfile generation fails, you can disable lockfiles by setting "
             "`[tool].lockfile = '<none>'`, e.g. setting `[black].lockfile`. You can also manually "
             "generate a lockfile, such as by using pip-compile or `pip freeze`. Set the "
-            "`[tool].lockfile` option to the path you manually generated. When manually maintaining "
-            "lockfiles, set `[python].invalid_lockfile_behavior = 'ignore'."
+            "`[tool].lockfile` option to the path you manually generated. When manually "
+            "maintaining lockfiles, set `[python].invalid_lockfile_behavior = 'ignore'."
         )
 
     if python_repos.repos:
@@ -132,64 +134,87 @@ async def generate_lockfile(
     generate_lockfiles_subsystem: GenerateLockfilesSubsystem,
     _: MaybeWarnPythonRepos,
 ) -> GenerateLockfileResult:
-    pyproject_toml = create_pyproject_toml(req.requirements, req.interpreter_constraints).encode()
-    pyproject_toml_digest, launcher_digest = await MultiGet(
-        Get(Digest, CreateDigest([FileContent("pyproject.toml", pyproject_toml)])),
-        Get(Digest, CreateDigest([POETRY_LAUNCHER])),
-    )
+    if req.use_pex:
+        result = await Get(
+            ProcessResult,
+            PexCliProcess(
+                subcommand=("lock", "create"),
+                # TODO: Hook up to [python-repos]. Don't call `MaybeWarnPythonRepos` in this case.
+                extra_args=(
+                    "--output=lock.json",
+                    # See https://github.com/pantsbuild/pants/issues/12458. For now, we always
+                    # generate universal locks because they have the best compatibility. We may
+                    # want to let users change this, as `style=strict` is safer.
+                    "--style=universal",
+                    "--resolver-version",
+                    "pip-2020-resolver",
+                    *req.interpreter_constraints.generate_pex_arg_list(),
+                    *req.requirements,
+                ),
+                output_files=("lock.json",),
+                description=req._description or f"Generate lockfile for {req.resolve_name}",
+                # Instead of caching lockfile generation with LMDB, we instead use the invalidation
+                # scheme from `lockfile_metadata.py` to check for stale/invalid lockfiles. This is
+                # necessary so that our invalidation is resilient to deleting LMDB or running on a
+                # new machine.
+                #
+                # We disable caching with LMDB so that when you generate a lockfile, you always get
+                # the most up-to-date snapshot of the world. This is generally desirable and also
+                # necessary to avoid an awkward edge case where different developers generate
+                # different lockfiles even when generating at the same time. See
+                # https://github.com/pantsbuild/pants/issues/12591.
+                cache_scope=ProcessCacheScope.PER_SESSION,
+            ),
+        )
+    else:
+        _pyproject_toml = create_pyproject_toml(
+            req.requirements, req.interpreter_constraints
+        ).encode()
+        _pyproject_toml_digest, _launcher_digest = await MultiGet(
+            Get(Digest, CreateDigest([FileContent("pyproject.toml", _pyproject_toml)])),
+            Get(Digest, CreateDigest([POETRY_LAUNCHER])),
+        )
 
-    poetry_pex = await Get(
-        VenvPex,
-        PexRequest(
-            output_filename="poetry.pex",
-            internal_only=True,
-            requirements=poetry_subsystem.pex_requirements(),
-            interpreter_constraints=poetry_subsystem.interpreter_constraints,
-            main=EntryPoint(PurePath(POETRY_LAUNCHER.path).stem),
-            sources=launcher_digest,
-        ),
-    )
+        _poetry_pex = await Get(
+            VenvPex,
+            PexRequest(
+                output_filename="poetry.pex",
+                internal_only=True,
+                requirements=poetry_subsystem.pex_requirements(),
+                interpreter_constraints=poetry_subsystem.interpreter_constraints,
+                main=EntryPoint(PurePath(POETRY_LAUNCHER.path).stem),
+                sources=_launcher_digest,
+            ),
+        )
 
-    # WONTFIX(#12314): Wire up Poetry to named_caches.
-    # WONTFIX(#12314): Wire up all the pip options like indexes.
-    poetry_lock_result = await Get(
-        ProcessResult,
-        VenvPexProcess(
-            poetry_pex,
-            argv=("lock",),
-            input_digest=pyproject_toml_digest,
-            output_files=("poetry.lock", "pyproject.toml"),
-            description=req._description or f"Generate lockfile for {req.resolve_name}",
-            # Instead of caching lockfile generation with LMDB, we instead use the invalidation
-            # scheme from `lockfile_metadata.py` to check for stale/invalid lockfiles. This is
-            # necessary so that our invalidation is resilient to deleting LMDB or running on a
-            # new machine.
-            #
-            # We disable caching with LMDB so that when you generate a lockfile, you always get
-            # the most up-to-date snapshot of the world. This is generally desirable and also
-            # necessary to avoid an awkward edge case where different developers generate different
-            # lockfiles even when generating at the same time. See
-            # https://github.com/pantsbuild/pants/issues/12591.
-            cache_scope=ProcessCacheScope.PER_SESSION,
-        ),
-    )
-    poetry_export_result = await Get(
-        ProcessResult,
-        VenvPexProcess(
-            poetry_pex,
-            argv=("export", "-o", req.lockfile_dest),
-            input_digest=poetry_lock_result.output_digest,
-            output_files=(req.lockfile_dest,),
-            description=(
-                f"Exporting Poetry lockfile to requirements.txt format for {req.resolve_name}"
+        # WONTFIX(#12314): Wire up Poetry to named_caches.
+        # WONTFIX(#12314): Wire up all the pip options like indexes.
+        _lock_result = await Get(
+            ProcessResult,
+            VenvPexProcess(
+                _poetry_pex,
+                argv=("lock",),
+                input_digest=_pyproject_toml_digest,
+                output_files=("poetry.lock", "pyproject.toml"),
+                description=req._description or f"Generate lockfile for {req.resolve_name}",
+                cache_scope=ProcessCacheScope.PER_SESSION,
             ),
-            level=LogLevel.DEBUG,
-        ),
-    )
+        )
+        result = await Get(
+            ProcessResult,
+            VenvPexProcess(
+                _poetry_pex,
+                argv=("export", "-o", req.lockfile_dest),
+                input_digest=_lock_result.output_digest,
+                output_files=(req.lockfile_dest,),
+                description=(
+                    f"Exporting Poetry lockfile to requirements.txt format for {req.resolve_name}"
+                ),
+                level=LogLevel.DEBUG,
+            ),
+        )
 
-    initial_lockfile_digest_contents = await Get(
-        DigestContents, Digest, poetry_export_result.output_digest
-    )
+    initial_lockfile_digest_contents = await Get(DigestContents, Digest, result.output_digest)
     # TODO(#12314) Improve error message on `Requirement.parse`
     metadata = PythonLockfileMetadata.new(
         req.interpreter_constraints,

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

@@ -3,20 +3,76 @@
 
 from __future__ import annotations
 
+import json
 from textwrap import dedent
 
 from pants.backend.python.goals.lockfile import (
     GeneratePythonLockfile,
     RequestedPythonUserResolveNames,
-    setup_user_lockfile_requests,
 )
+from pants.backend.python.goals.lockfile import rules as lockfile_rules
+from pants.backend.python.goals.lockfile import setup_user_lockfile_requests
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import PythonRequirementTarget
+from pants.backend.python.util_rules import pex
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
-from pants.core.goals.generate_lockfiles import UserGenerateLockfiles
+from pants.core.goals.generate_lockfiles import GenerateLockfileResult, UserGenerateLockfiles
+from pants.engine.fs import DigestContents
 from pants.engine.rules import SubsystemRule
 from pants.testutil.rule_runner import PYTHON_BOOTSTRAP_ENV, QueryRule, RuleRunner
 from pants.util.ordered_set import FrozenOrderedSet
+from pants.util.strutil import strip_prefix
+
+
+def test_lockfile_generation() -> None:
+    rule_runner = RuleRunner(
+        rules=[
+            *lockfile_rules(),
+            *pex.rules(),
+            QueryRule(GenerateLockfileResult, [GeneratePythonLockfile]),
+        ]
+    )
+    rule_runner.set_options([], env_inherit=PYTHON_BOOTSTRAP_ENV)
+    result = rule_runner.request(
+        GenerateLockfileResult,
+        [
+            GeneratePythonLockfile(
+                requirements=FrozenOrderedSet(["ansicolors==1.1.8"]),
+                interpreter_constraints=InterpreterConstraints(),
+                resolve_name="test",
+                lockfile_dest="test.lock",
+                use_pex=True,
+            )
+        ],
+    )
+
+    digest_contents = rule_runner.request(DigestContents, [result.digest])
+    assert len(digest_contents) == 1
+    lock_content = digest_contents[0].content.decode()
+    header = dedent(
+        """\
+        # This lockfile was autogenerated by Pants. To regenerate, run:
+        #
+        #    ./pants generate-lockfiles --resolve=test
+        #
+        # --- BEGIN PANTS LOCKFILE METADATA: DO NOT EDIT OR REMOVE ---
+        # {
+        #   "version": 2,
+        #   "valid_for_interpreter_constraints": [],
+        #   "generated_with_requirements": [
+        #     "ansicolors==1.1.8"
+        #   ]
+        # }
+        # --- END PANTS LOCKFILE METADATA ---
+        """
+    )
+    assert lock_content.startswith(header)
+
+    lock_entry = json.loads(strip_prefix(lock_content, header))
+    reqs = lock_entry["locked_resolves"][0]["locked_requirements"]
+    assert len(reqs) == 1
+    assert reqs[0]["project_name"] == "ansicolors"
+    assert reqs[0]["version"] == "1.1.8"
 
 
 def test_multiple_resolves() -> None: