[internal] Add experimental per-tool lockfiles with Docformatter as an example

pants.toml

@@ -77,6 +77,7 @@ interpreter_constraints = [">=3.7,<3.10"]
 
 [docformatter]
 args = ["--wrap-summaries=100", "--wrap-descriptions=100"]
+experimental_lockfile = "src/python/pants/backend/python/lint/docformatter/lockfile.txt"
 
 [flake8]
 config = "build-support/flake8/.flake8"

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

@@ -4,18 +4,20 @@
 from __future__ import annotations
 
 import logging
+from dataclasses import dataclass
 from typing import cast
 
 from pants.backend.python.subsystems.python_tool_base import PythonToolBase
 from pants.backend.python.target_types import ConsoleScript, PythonRequirementsField
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.pex import PexRequest, PexRequirements, VenvPex, VenvPexProcess
 from pants.engine.addresses import Addresses
-from pants.engine.fs import CreateDigest, Digest, FileContent, Workspace
+from pants.engine.fs import CreateDigest, Digest, FileContent, MergeDigests, Workspace
 from pants.engine.goal import Goal, GoalSubsystem
 from pants.engine.process import ProcessResult
-from pants.engine.rules import Get, MultiGet, collect_rules, goal_rule
+from pants.engine.rules import Get, MultiGet, collect_rules, goal_rule, rule
 from pants.engine.target import TransitiveTargets, TransitiveTargetsRequest
+from pants.engine.unions import UnionMembership, union
 from pants.python.python_setup import PythonSetup
 from pants.util.strutil import pluralize
 
@@ -46,6 +48,11 @@ def lockfile_dest(self) -> str:
         return cast(str, self.options.lockfile_dest)
 
 
+# --------------------------------------------------------------------------------------
+# User lockfiles
+# --------------------------------------------------------------------------------------
+
+
 class LockSubsystem(GoalSubsystem):
     name = "lock"
     help = "Generate a lockfile."
@@ -104,8 +111,12 @@ async def generate_lockfile(
         return LockGoal(exit_code=0)
 
     input_requirements_get = Get(
-        Digest, CreateDigest([FileContent("requirements.in", "\n".join(reqs).encode())])
+        Digest, CreateDigest([FileContent("requirements.in", "\n".join(reqs.req_strings).encode())])
     )
+    # TODO: Figure out which interpreter constraints to use...Likely get it from the
+    #  transitive closure. When we're doing a single global lockfile, it's fine to do that,
+    #  but we need to figure out how this will work with multiple resolves.
+    interpreter_constraints = InterpreterConstraints(python_setup.interpreter_constraints)
     # TODO(#12314): Figure out named_caches for pip-tools. The best would be to share the cache
     #  between Pex and Pip. Next best is a dedicated named_cache.
     pip_compile_get = Get(
@@ -114,11 +125,11 @@ async def generate_lockfile(
             output_filename="pip_compile.pex",
             internal_only=True,
             requirements=PexRequirements(pip_tools_subsystem.all_requirements),
-            # TODO: Figure out which interpreter constraints to use...Likely get it from the
-            #  transitive closure. When we're doing a single global lockfile, it's fine to do that,
-            #  but we need to figure out how this will work with multiple resolves.
-            interpreter_constraints=InterpreterConstraints(python_setup.interpreter_constraints),
+            interpreter_constraints=interpreter_constraints,
             main=pip_tools_subsystem.main,
+            description=(
+                f"Building pip_compile.pex with interpreter constraints: {interpreter_constraints}"
+            ),
         ),
     )
     input_requirements, pip_compile = await MultiGet(input_requirements_get, pip_compile_get)
@@ -129,7 +140,8 @@ async def generate_lockfile(
         VenvPexProcess(
             pip_compile,
             description=(
-                f"Generate lockfile for {pluralize(len(reqs), 'requirements')}: {', '.join(reqs)}"
+                f"Generate lockfile for {pluralize(len(reqs.req_strings), 'requirements')}: "
+                f"{', '.join(reqs.req_strings)}"
             ),
             # TODO(#12314): Wire up all the pip options like indexes.
             argv=[
@@ -152,5 +164,120 @@ async def generate_lockfile(
     return LockGoal(exit_code=0)
 
 
+# --------------------------------------------------------------------------------------
+# Tool lockfiles
+# --------------------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class PythonToolLockfile:
+    digest: Digest
+    tool_name: str
+    path: str
+
+
+@union
+class PythonToolLockfileSentinel:
+    pass
+
+
+@dataclass(frozen=True)
+class PythonToolLockfileRequest:
+    tool_name: str
+    lockfile_path: str
+    requirements: tuple[str, ...]
+    interpreter_constraints: tuple[str, ...]
+
+
+# TODO(#12314): Unify this goal with `lock` once we figure out how to unify the semantics,
+#  particularly w/ CLI specs. This is a separate goal only to facilitate progress.
+class ToolLockSubsystem(GoalSubsystem):
+    name = "tool-lock"
+    help = "Generate a lockfile for a Python tool."
+    required_union_implementations = (PythonToolLockfileSentinel,)
+
+
+class ToolLockGoal(Goal):
+    subsystem_cls = ToolLockSubsystem
+
+
+@rule
+async def generate_tool_lockfile(
+    request: PythonToolLockfileRequest, pip_tools_subsystem: PipToolsSubsystem
+) -> PythonToolLockfile:
+    input_requirements_get = Get(
+        Digest,
+        CreateDigest(
+            [
+                FileContent(
+                    f"requirements_{request.tool_name}.in", "\n".join(request.requirements).encode()
+                )
+            ]
+        ),
+    )
+    interpreter_constraints = InterpreterConstraints(request.interpreter_constraints)
+    pip_compile_get = Get(
+        VenvPex,
+        PexRequest(
+            output_filename="pip_compile.pex",
+            internal_only=True,
+            requirements=PexRequirements(pip_tools_subsystem.all_requirements),
+            interpreter_constraints=interpreter_constraints,
+            main=pip_tools_subsystem.main,
+            description=(
+                f"Building pip_compile.pex with interpreter constraints: {interpreter_constraints}"
+            ),
+        ),
+    )
+    input_requirements, pip_compile = await MultiGet(input_requirements_get, pip_compile_get)
+
+    result = await Get(
+        ProcessResult,
+        VenvPexProcess(
+            pip_compile,
+            description=f"Generate lockfile for {request.tool_name}",
+            argv=[
+                f"requirements_{request.tool_name}.in",
+                "--generate-hashes",
+                f"--output-file={request.lockfile_path}",
+                # NB: This allows pinning setuptools et al, which we must do. This will become
+                # the default in a future version of pip-tools.
+                "--allow-unsafe",
+            ],
+            input_digest=input_requirements,
+            output_files=(request.lockfile_path,),
+        ),
+    )
+
+    return PythonToolLockfile(result.output_digest, request.tool_name, request.lockfile_path)
+
+
+@goal_rule
+async def generate_all_tool_lockfiles(
+    workspace: Workspace,
+    union_membership: UnionMembership,
+) -> ToolLockGoal:
+    # TODO(#12314): Add logic to inspect the Specs and generate for only relevant lockfiles. For
+    #  now, we generate for all tools.
+    requests = await MultiGet(
+        Get(PythonToolLockfileRequest, PythonToolLockfileSentinel, sentinel())
+        for sentinel in union_membership.get(PythonToolLockfileSentinel)
+    )
+    if not requests:
+        return ToolLockGoal(exit_code=0)
+
+    results = await MultiGet(
+        Get(PythonToolLockfile, PythonToolLockfileRequest, req)
+        for req in requests
+        if req.lockfile_path not in {"<none>", "<default>"}
+    )
+    merged_digest = await Get(Digest, MergeDigests(res.digest for res in results))
+    workspace.write_digest(merged_digest)
+    for result in results:
+        logger.info(f"Wrote lockfile for {result.tool_name} to {result.path}")
+
+    return ToolLockGoal(exit_code=0)
+
+
 def rules():
     return collect_rules()

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

@@ -695,7 +695,7 @@ async def get_requirements(
         for tgt in direct_deps_with_excl
         if tgt.has_field(PythonRequirementsField)
     )
-    req_strs = list(reqs)
+    req_strs = list(reqs.req_strings)
 
     # Add the requirements on any exported targets on which we depend.
     kwargs_for_exported_targets_we_depend_on = await MultiGet(

src/python/pants/backend/python/lint/docformatter/BUILD

@@ -1,5 +1,6 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-python_library()
+python_library(dependencies=[":lockfile"])
 python_tests(name="tests", timeout=120)
+resources(name="lockfile", sources=["lockfile.txt"])

src/python/pants/backend/python/lint/docformatter/lockfile.txt

@@ -0,0 +1,12 @@
+#
+# This file is autogenerated by pip-compile with python 3.6
+# To update, run:
+#
+#    pip-compile --allow-unsafe --generate-hashes --output-file=src/python/pants/backend/python/lint/docformatter/lockfile.txt requirements_docformatter.in
+#
+docformatter==1.4 \
+    --hash=sha256:064e6d81f04ac96bc0d176cbaae953a0332482b22d3ad70d47c8a7f2732eef6f
+    # via -r requirements_docformatter.in
+untokenize==0.1.1 \
+    --hash=sha256:3865dbbbb8efb4bb5eaa72f1be7f3e0be00ea8b7f125c69cbd1f5fda926f37a2
+    # via docformatter

src/python/pants/backend/python/lint/docformatter/register.py

@@ -8,9 +8,9 @@
 """
 
 from pants.backend.python.lint import python_fmt
-from pants.backend.python.lint.docformatter import skip_field
+from pants.backend.python.lint.docformatter import skip_field, subsystem
 from pants.backend.python.lint.docformatter.rules import rules as docformatter_rules
 
 
 def rules():
-    return (*docformatter_rules(), *python_fmt.rules(), *skip_field.rules())
+    return (*docformatter_rules(), *python_fmt.rules(), *skip_field.rules(), *subsystem.rules())

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

@@ -1,6 +1,7 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+import importlib.resources
 from dataclasses import dataclass
 from typing import Tuple
 
@@ -14,7 +15,7 @@
 from pants.core.goals.fmt import FmtResult
 from pants.core.goals.lint import LintRequest, LintResult, LintResults
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
-from pants.engine.fs import Digest
+from pants.engine.fs import Digest, FileContent
 from pants.engine.process import FallibleProcessResult, Process, ProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import FieldSet, Target
@@ -58,23 +59,40 @@ def generate_args(
 
 @rule(level=LogLevel.DEBUG)
 async def setup_docformatter(setup_request: SetupRequest, docformatter: Docformatter) -> Setup:
-    docformatter_pex_request = Get(
+    if docformatter.lockfile == "<none>":
+        requirements = PexRequirements(docformatter.all_requirements)
+    elif docformatter.lockfile == "<default>":
+        requirements = PexRequirements(
+            file_content=FileContent(
+                "docformatter_default_lockfile.txt",
+                importlib.resources.read_binary(
+                    "pants.backend.python.lint.docformatter", "lockfile.txt"
+                ),
+            )
+        )
+    else:
+        requirements = PexRequirements(
+            file_path=docformatter.lockfile,
+            file_path_description_of_origin="the option `[docformatter].experimental_lockfile`",
+        )
+
+    docformatter_pex_get = Get(
         VenvPex,
         PexRequest(
             output_filename="docformatter.pex",
             internal_only=True,
-            requirements=PexRequirements(docformatter.all_requirements),
             interpreter_constraints=InterpreterConstraints(docformatter.interpreter_constraints),
             main=docformatter.main,
+            description="Build docformatter.pex",
+            requirements=requirements,
+            is_lockfile=docformatter.lockfile != "<none>",
         ),
     )
-
-    source_files_request = Get(
+    source_files_get = Get(
         SourceFiles,
         SourceFilesRequest(field_set.sources for field_set in setup_request.request.field_sets),
     )
-
-    source_files, docformatter_pex = await MultiGet(source_files_request, docformatter_pex_request)
+    source_files, docformatter_pex = await MultiGet(source_files_get, docformatter_pex_get)
 
     source_files_snapshot = (
         source_files.snapshot

src/python/pants/backend/python/lint/docformatter/subsystem.py

@@ -3,8 +3,14 @@
 
 from typing import Tuple, cast
 
+from pants.backend.experimental.python.lockfile import (
+    PythonToolLockfileRequest,
+    PythonToolLockfileSentinel,
+)
 from pants.backend.python.subsystems.python_tool_base import PythonToolBase
 from pants.backend.python.target_types import ConsoleScript
+from pants.engine.rules import collect_rules, rule
+from pants.engine.unions import UnionRule
 from pants.option.custom_types import shell_str
 
 
@@ -38,6 +44,26 @@ def register_options(cls, register):
                 f'`--{cls.options_scope}-args="--wrap-summaries=100 --pre-summary-newline"`.'
             ),
         )
+        register(
+            "--experimental-lockfile",
+            type=str,
+            default="<none>",
+            advanced=True,
+            help=(
+                "Path to a lockfile used for installing the tool.\n\n"
+                "Set to the string '<default>' to use a lockfile provided by "
+                "Pants, so long as you have not changed the `--version`, `--extra-requirements`, "
+                "and `--interpreter-constraints` options. See {} for the default lockfile "
+                "contents.\n\n"
+                "Set to the string '<none>' to opt out of using a lockfile. We do not recommend "
+                "this, as lockfiles are essential for reproducible builds.\n\n"
+                "To use a custom lockfile, set this option to a file path relative to the build "
+                "root, then activate the backend_package `pants.backend.experimental.python` and "
+                "run `./pants tool-lock`.\n\n"
+                "This option is experimental and will likely change. It does not follow the normal "
+                "deprecation cycle."
+            ),
+        )
 
     @property
     def skip(self) -> bool:
@@ -46,3 +72,27 @@ def skip(self) -> bool:
     @property
     def args(self) -> Tuple[str, ...]:
         return tuple(self.options.args)
+
+    @property
+    def lockfile(self) -> str:
+        return cast(str, self.options.experimental_lockfile)
+
+
+class DocformatterLockfileSentinel(PythonToolLockfileSentinel):
+    pass
+
+
+@rule
+def setup_lockfile_request(
+    _: DocformatterLockfileSentinel, docformatter: Docformatter
+) -> PythonToolLockfileRequest:
+    return PythonToolLockfileRequest(
+        tool_name=docformatter.options_scope,
+        lockfile_path=docformatter.lockfile,
+        requirements=docformatter.all_requirements,
+        interpreter_constraints=docformatter.interpreter_constraints,
+    )
+
+
+def rules():
+    return (*collect_rules(), UnionRule(PythonToolLockfileSentinel, DocformatterLockfileSentinel))

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

@@ -133,7 +133,9 @@ async def pylint_lint_partition(partition: PylintPartition, pylint: Pylint) -> L
         PexRequest(
             output_filename="pylint.pex",
             internal_only=True,
-            requirements=PexRequirements([*pylint.all_requirements, *plugin_requirements]),
+            requirements=PexRequirements(
+                [*pylint.all_requirements, *plugin_requirements.req_strings]
+            ),
             interpreter_constraints=partition.interpreter_constraints,
         ),
     )