Add Protobuf formatting using buf format

src/python/pants/backend/codegen/protobuf/lint/buf/format_rules.py

@@ -0,0 +1,157 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+from dataclasses import dataclass
+
+from pants.backend.codegen.protobuf.lint.buf.skip_field import SkipBufField
+from pants.backend.codegen.protobuf.lint.buf.subsystem import BufSubsystem
+from pants.backend.codegen.protobuf.target_types import (
+    ProtobufDependenciesField,
+    ProtobufSourceField,
+)
+from pants.core.goals.fmt import FmtRequest, FmtResult
+from pants.core.goals.lint import LintResult, LintResults, LintTargetsRequest
+from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
+from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
+from pants.core.util_rules.system_binaries import (
+    BinaryShims,
+    BinaryShimsRequest,
+    DiffBinary,
+    DiffBinaryRequest,
+)
+from pants.engine.fs import Digest, MergeDigests
+from pants.engine.internals.native_engine import Snapshot
+from pants.engine.platform import Platform
+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
+from pants.engine.unions import UnionRule
+from pants.util.logging import LogLevel
+from pants.util.strutil import pluralize
+
+
+@dataclass(frozen=True)
+class BufFieldSet(FieldSet):
+    required_fields = (ProtobufSourceField,)
+
+    sources: ProtobufSourceField
+    dependencies: ProtobufDependenciesField
+
+    @classmethod
+    def opt_out(cls, tgt: Target) -> bool:
+        return tgt.get(SkipBufField).value
+
+
+class BufFormatRequest(LintTargetsRequest, FmtRequest):
+    field_set_type = BufFieldSet
+    name = "buf-format"
+
+
+@dataclass(frozen=True)
+class SetupRequest:
+    request: BufFormatRequest
+    check_only: bool
+
+
+@dataclass(frozen=True)
+class Setup:
+    process: Process
+    original_snapshot: Snapshot
+
+
+@rule(level=LogLevel.DEBUG)
+async def setup_buf_format(setup_request: SetupRequest, buf: BufSubsystem) -> Setup:
+    diff_binary = await Get(DiffBinary, DiffBinaryRequest())
+    download_buf_get = Get(
+        DownloadedExternalTool, ExternalToolRequest, buf.get_request(Platform.current)
+    )
+    binary_shims_get = Get(
+        BinaryShims,
+        BinaryShimsRequest,
+        BinaryShimsRequest.for_paths(
+            diff_binary,
+            rationale="buf format requires diff in linting mode",
+            output_directory=".bin",
+        ),
+    )
+    source_files_get = Get(
+        SourceFiles,
+        SourceFilesRequest(field_set.sources for field_set in setup_request.request.field_sets),
+    )
+    downloaded_buf, binary_shims, source_files = await MultiGet(
+        download_buf_get, binary_shims_get, source_files_get
+    )
+
+    source_files_snapshot = (
+        source_files.snapshot
+        if setup_request.request.prior_formatter_result is None
+        else setup_request.request.prior_formatter_result
+    )
+
+    input_digest = await Get(
+        Digest,
+        MergeDigests((source_files_snapshot.digest, downloaded_buf.digest, binary_shims.digest)),
+    )
+
+    argv = [
+        downloaded_buf.exe,
+        "format",
+        # If linting, use `-d` to error with a diff. Else, write the change with `-w`.
+        *(["-d"] if setup_request.check_only else ["-w"]),
+        *buf.args,
+        "--path",
+        ",".join(source_files_snapshot.files),
+    ]
+    process = Process(
+        argv=argv,
+        input_digest=input_digest,
+        output_files=source_files_snapshot.files,
+        description=f"Run buf format on {pluralize(len(setup_request.request.field_sets), 'file')}.",
+        level=LogLevel.DEBUG,
+        env={
+            "PATH": binary_shims.bin_directory,
+        },
+    )
+    return Setup(process, original_snapshot=source_files_snapshot)
+
+
+@rule(desc="Format with buf format", level=LogLevel.DEBUG)
+async def run_buf_format(request: BufFormatRequest, buf: BufSubsystem) -> FmtResult:
+    if buf.skip:
+        return FmtResult.skip(formatter_name=request.name)
+    setup = await Get(Setup, SetupRequest(request, check_only=False))
+    result = await Get(ProcessResult, Process, setup.process)
+    output_snapshot = await Get(Snapshot, Digest, result.output_digest)
+    return FmtResult(
+        setup.original_snapshot,
+        output_snapshot,
+        stdout=result.stdout.decode(),
+        stderr=result.stderr.decode(),
+        formatter_name=request.name,
+    )
+
+
+@rule(desc="Lint with buf format", level=LogLevel.DEBUG)
+async def run_buf_lint(request: BufFormatRequest, buf: BufSubsystem) -> LintResults:
+    if buf.skip:
+        return LintResults([], linter_name=request.name)
+    setup = await Get(Setup, SetupRequest(request, check_only=True))
+    result = await Get(FallibleProcessResult, Process, setup.process)
+
+    return LintResults(
+        [
+            LintResult(
+                exit_code=0 if not result.stdout else 1,  # buf format always exits with code 0
+                stdout=result.stdout.decode(),
+                stderr=result.stderr.decode(),
+            )
+        ],
+        linter_name=request.name,
+    )
+
+
+def rules():
+    return [
+        *collect_rules(),
+        UnionRule(FmtRequest, BufFormatRequest),
+        UnionRule(LintTargetsRequest, BufFormatRequest),
+    ]

src/python/pants/backend/codegen/protobuf/lint/buf/format_rules_integration_test.py

@@ -0,0 +1,143 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+import pytest
+
+from pants.backend.codegen.protobuf.lint.buf.format_rules import BufFieldSet, BufFormatRequest
+from pants.backend.codegen.protobuf.lint.buf.format_rules import rules as buf_rules
+from pants.backend.codegen.protobuf.target_types import ProtobufSourcesGeneratorTarget
+from pants.backend.codegen.protobuf.target_types import rules as target_types_rules
+from pants.core.goals.fmt import FmtResult
+from pants.core.goals.lint import LintResult, LintResults
+from pants.core.util_rules import config_files, external_tool, source_files
+from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
+from pants.engine.addresses import Address
+from pants.engine.fs import CreateDigest, FileContent
+from pants.engine.internals.native_engine import Digest, Snapshot
+from pants.engine.target import Target
+from pants.testutil.rule_runner import QueryRule, RuleRunner
+
+
+@pytest.fixture
+def rule_runner() -> RuleRunner:
+    return RuleRunner(
+        rules=[
+            *buf_rules(),
+            *config_files.rules(),
+            *external_tool.rules(),
+            *source_files.rules(),
+            *target_types_rules(),
+            QueryRule(LintResults, [BufFormatRequest]),
+            QueryRule(FmtResult, [BufFormatRequest]),
+            QueryRule(SourceFiles, [SourceFilesRequest]),
+        ],
+        target_types=[ProtobufSourcesGeneratorTarget],
+    )
+
+
+GOOD_FILE = 'syntax = "proto3";\n\npackage foo.v1;\nmessage Foo {}\n'
+BAD_FILE = 'syntax = "proto3";\n\npackage foo.v1;\nmessage Foo {\n}\n'
+
+
+def run_buf(
+    rule_runner: RuleRunner,
+    targets: list[Target],
+    *,
+    extra_args: list[str] | None = None,
+) -> tuple[tuple[LintResult, ...], FmtResult]:
+    rule_runner.set_options(
+        [
+            "--backend-packages=pants.backend.codegen.protobuf.lint.buf",
+            *(extra_args or ()),
+        ],
+        env_inherit={"PATH"},
+    )
+    field_sets = [BufFieldSet.create(tgt) for tgt in targets]
+    results = rule_runner.request(LintResults, [BufFormatRequest(field_sets)])
+    input_sources = rule_runner.request(
+        SourceFiles,
+        [
+            SourceFilesRequest(field_set.sources for field_set in field_sets),
+        ],
+    )
+    fmt_result = rule_runner.request(
+        FmtResult,
+        [
+            BufFormatRequest(field_sets, prior_formatter_result=input_sources.snapshot),
+        ],
+    )
+
+    return results.results, fmt_result
+
+
+def get_snapshot(rule_runner: RuleRunner, source_files: dict[str, str]) -> Snapshot:
+    files = [FileContent(path, content.encode()) for path, content in source_files.items()]
+    digest = rule_runner.request(Digest, [CreateDigest(files)])
+    return rule_runner.request(Snapshot, [digest])
+
+
+def test_passing(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files({"f.proto": GOOD_FILE, "BUILD": "protobuf_sources(name='t')"})
+    tgt = rule_runner.get_target(Address("", target_name="t", relative_file_path="f.proto"))
+    lint_results, fmt_result = run_buf(rule_runner, [tgt])
+    assert len(lint_results) == 1
+    assert lint_results[0].exit_code == 0
+    assert lint_results[0].stdout == ""
+    assert lint_results[0].stderr == ""
+    assert fmt_result.stdout == ""
+    assert fmt_result.output == get_snapshot(rule_runner, {"f.proto": GOOD_FILE})
+    assert fmt_result.did_change is False
+
+
+def test_failing(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files({"f.proto": BAD_FILE, "BUILD": "protobuf_sources(name='t')"})
+    tgt = rule_runner.get_target(Address("", target_name="t", relative_file_path="f.proto"))
+    lint_results, fmt_result = run_buf(rule_runner, [tgt])
+    assert len(lint_results) == 1
+    assert lint_results[0].exit_code == 1
+    assert "f.proto.orig" in lint_results[0].stdout
+    assert fmt_result.output == get_snapshot(rule_runner, {"f.proto": GOOD_FILE})
+    assert fmt_result.did_change is True
+
+
+def test_multiple_targets(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {"good.proto": GOOD_FILE, "bad.proto": BAD_FILE, "BUILD": "protobuf_sources(name='t')"}
+    )
+    tgts = [
+        rule_runner.get_target(Address("", target_name="t", relative_file_path="good.proto")),
+        rule_runner.get_target(Address("", target_name="t", relative_file_path="bad.proto")),
+    ]
+    lint_results, fmt_result = run_buf(rule_runner, tgts)
+    assert len(lint_results) == 1
+    assert lint_results[0].exit_code == 1
+    assert "bad.proto.orig" in lint_results[0].stdout
+    assert "good.proto" not in lint_results[0].stdout
+    assert fmt_result.output == get_snapshot(
+        rule_runner, {"good.proto": GOOD_FILE, "bad.proto": GOOD_FILE}
+    )
+    assert fmt_result.did_change is True
+
+
+def test_passthrough_args(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files({"f.proto": GOOD_FILE, "BUILD": "protobuf_sources(name='t')"})
+    tgt = rule_runner.get_target(Address("", target_name="t", relative_file_path="f.proto"))
+    lint_results, fmt_result = run_buf(rule_runner, [tgt], extra_args=["--buf-args=--debug"])
+    assert len(lint_results) == 1
+    assert lint_results[0].exit_code == 0
+    assert lint_results[0].stdout == ""
+    assert "DEBUG" in lint_results[0].stderr
+    assert fmt_result.stdout == ""
+    assert fmt_result.output == get_snapshot(rule_runner, {"f.proto": GOOD_FILE})
+    assert fmt_result.did_change is False
+
+
+def test_skip(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files({"f.proto": BAD_FILE, "BUILD": "protobuf_sources(name='t')"})
+    tgt = rule_runner.get_target(Address("", target_name="t", relative_file_path="f.proto"))
+    lint_results, fmt_result = run_buf(rule_runner, [tgt], extra_args=["--buf-skip"])
+    assert not lint_results
+    assert fmt_result.skipped is True
+    assert fmt_result.did_change is False

src/python/pants/backend/codegen/protobuf/lint/buf/rules.py → src/python/pants/backend/codegen/protobuf/lint/buf/lint_rules.py

@@ -34,13 +34,13 @@ def opt_out(cls, tgt: Target) -> bool:
         return tgt.get(SkipBufField).value
 
 
-class BufRequest(LintTargetsRequest):
+class BufLintRequest(LintTargetsRequest):
     field_set_type = BufFieldSet
-    name = BufSubsystem.options_scope
+    name = "buf-lint"
 
 
-@rule(desc="Lint with Buf", level=LogLevel.DEBUG)
-async def run_buf(request: BufRequest, buf: BufSubsystem) -> LintResults:
+@rule(desc="Lint with buf lint", level=LogLevel.DEBUG)
+async def run_buf(request: BufLintRequest, buf: BufSubsystem) -> LintResults:
     if buf.skip:
         return LintResults([], linter_name=request.name)
 
@@ -96,7 +96,7 @@ async def run_buf(request: BufRequest, buf: BufSubsystem) -> LintResults:
                 ",".join(target_sources_stripped.snapshot.files),
             ],
             input_digest=input_digest,
-            description=f"Run Buf on {pluralize(len(request.field_sets), 'file')}.",
+            description=f"Run buf lint on {pluralize(len(request.field_sets), 'file')}.",
             level=LogLevel.DEBUG,
         ),
     )
@@ -106,4 +106,4 @@ async def run_buf(request: BufRequest, buf: BufSubsystem) -> LintResults:
 
 
 def rules():
-    return [*collect_rules(), UnionRule(LintTargetsRequest, BufRequest)]
+    return [*collect_rules(), UnionRule(LintTargetsRequest, BufLintRequest)]

src/python/pants/backend/codegen/protobuf/lint/buf/rules_integration_test.py → src/python/pants/backend/codegen/protobuf/lint/buf/lint_rules_integration_test.py

@@ -8,8 +8,8 @@
 
 import pytest
 
-from pants.backend.codegen.protobuf.lint.buf.rules import BufFieldSet, BufRequest
-from pants.backend.codegen.protobuf.lint.buf.rules import rules as buf_rules
+from pants.backend.codegen.protobuf.lint.buf.lint_rules import BufFieldSet, BufLintRequest
+from pants.backend.codegen.protobuf.lint.buf.lint_rules import rules as buf_rules
 from pants.backend.codegen.protobuf.target_types import ProtobufSourcesGeneratorTarget
 from pants.backend.codegen.protobuf.target_types import rules as target_types_rules
 from pants.core.goals.lint import LintResult, LintResults
@@ -28,7 +28,7 @@ def rule_runner() -> RuleRunner:
             *external_tool.rules(),
             *stripped_source_files.rules(),
             *target_types_rules(),
-            QueryRule(LintResults, [BufRequest]),
+            QueryRule(LintResults, [BufLintRequest]),
         ],
         target_types=[ProtobufSourcesGeneratorTarget],
     )
@@ -55,7 +55,7 @@ def run_buf(
     )
     results = rule_runner.request(
         LintResults,
-        [BufRequest(BufFieldSet.create(tgt) for tgt in targets)],
+        [BufLintRequest(BufFieldSet.create(tgt) for tgt in targets)],
     )
     return results.results
 
@@ -133,7 +133,7 @@ def test_passthrough_args(rule_runner: RuleRunner) -> None:
     assert_success(
         rule_runner,
         tgt,
-        extra_args=[f"--buf-lint-args='--config={config}'"],
+        extra_args=[f"--buf-args='--config={config}'"],
     )
 
 
@@ -142,7 +142,7 @@ def test_skip(rule_runner: RuleRunner) -> None:
         {"foo/v1/f.proto": BAD_FILE, "foo/v1/BUILD": "protobuf_sources(name='t')"}
     )
     tgt = rule_runner.get_target(Address("foo/v1", target_name="t", relative_file_path="f.proto"))
-    result = run_buf(rule_runner, [tgt], extra_args=["--buf-lint-skip"])
+    result = run_buf(rule_runner, [tgt], extra_args=["--buf-skip"])
     assert not result
 
 

src/python/pants/backend/codegen/protobuf/lint/buf/register.py

@@ -2,8 +2,9 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 from pants.backend.codegen.protobuf.lint.buf import skip_field
-from pants.backend.codegen.protobuf.lint.buf.rules import rules as buf_rules
+from pants.backend.codegen.protobuf.lint.buf.format_rules import rules as buf_format_rules
+from pants.backend.codegen.protobuf.lint.buf.lint_rules import rules as buf_lint_rules
 
 
 def rules():
-    return (*buf_rules(), *skip_field.rules())
+    return (*buf_format_rules(), *buf_lint_rules(), *skip_field.rules())