Prefactor tffmt to be friendlier to upcoming changes

src/python/pants/backend/experimental/terraform/register.py

@@ -3,7 +3,7 @@
 
 from pants.backend.python.goals import lockfile as python_lockfile
 from pants.backend.python.util_rules.pex import rules as pex_rules
-from pants.backend.terraform import dependency_inference, style, tool
+from pants.backend.terraform import dependency_inference, tool
 from pants.backend.terraform.goals import check, tailor
 from pants.backend.terraform.lint.tffmt.tffmt import rules as tffmt_rules
 from pants.backend.terraform.target_types import TerraformModuleTarget
@@ -23,7 +23,6 @@ def rules():
         *tailor.rules(),
         *target_types_rules(),
         *tool.rules(),
-        *style.rules(),
         *pex_rules(),
         *tffmt_rules(),
         *python_lockfile.rules(),

src/python/pants/backend/terraform/goals/check.py

@@ -2,16 +2,18 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 
-from pants.backend.terraform.style import StyleSetup, StyleSetupRequest
+from pants.backend.terraform.partition import partition_files_by_directory
 from pants.backend.terraform.target_types import TerraformFieldSet
 from pants.backend.terraform.tool import TerraformProcess
 from pants.core.goals.check import CheckRequest, CheckResult, CheckResults
+from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
 from pants.engine.internals.selectors import Get, MultiGet
 from pants.engine.process import FallibleProcessResult
 from pants.engine.rules import collect_rules, rule
 from pants.engine.unions import UnionRule
 from pants.option.option_types import SkipOption
 from pants.option.subsystem import Subsystem
+from pants.util.strutil import pluralize
 
 
 class TerraformValidateSubsystem(Subsystem):
@@ -34,14 +36,26 @@ async def terraform_check(
     if subsystem.skip:
         return CheckResults([], checker_name=request.name)
 
-    setup = await Get(StyleSetup, StyleSetupRequest(request, ("validate",)))
+    source_files = await Get(
+        SourceFiles, SourceFilesRequest([field_set.sources for field_set in request.field_sets])
+    )
+    files_by_directory = partition_files_by_directory(source_files.files)
+
     results = await MultiGet(
-        Get(FallibleProcessResult, TerraformProcess, process)
-        for _, (process, _) in setup.directory_to_process.items()
+        Get(
+            FallibleProcessResult,
+            TerraformProcess(
+                args=("validate", directory),
+                input_digest=source_files.snapshot.digest,
+                output_files=tuple(files),
+                description=f"Run `terraform fmt` on {pluralize(len(files), 'file')}.",
+            ),
+        )
+        for directory, files in files_by_directory.items()
     )
 
     check_results = []
-    for directory, result in zip(setup.directory_to_process.keys(), results):
+    for directory, result in zip(files_by_directory, results):
         check_results.append(
             CheckResult.from_fallible_process_result(
                 result, partition_description=f"`terraform validate` on `{directory}`"

src/python/pants/backend/terraform/goals/check_test.py

@@ -7,7 +7,7 @@
 
 import pytest
 
-from pants.backend.terraform import style, tool
+from pants.backend.terraform import tool
 from pants.backend.terraform.goals import check
 from pants.backend.terraform.goals.check import TerraformCheckRequest
 from pants.backend.terraform.target_types import TerraformFieldSet, TerraformModuleTarget
@@ -28,7 +28,6 @@ def rule_runner() -> RuleRunner:
             *external_tool.rules(),
             *check.rules(),
             *tool.rules(),
-            *style.rules(),
             *source_files.rules(),
             QueryRule(CheckResults, (TerraformCheckRequest,)),
             QueryRule(SourceFiles, (SourceFilesRequest,)),

src/python/pants/backend/terraform/lint/tffmt/tffmt.py

@@ -1,14 +1,19 @@
 # Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
+from __future__ import annotations
+
 import logging
 import textwrap
+from dataclasses import dataclass
+from typing import Any, cast
 
-from pants.backend.terraform.style import StyleSetup, StyleSetupRequest
+from pants.backend.terraform.partition import partition_files_by_directory
 from pants.backend.terraform.target_types import TerraformFieldSet
 from pants.backend.terraform.tool import TerraformProcess
 from pants.backend.terraform.tool import rules as tool_rules
 from pants.core.goals.fmt import FmtResult, FmtTargetsRequest
 from pants.core.util_rules import external_tool
+from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
 from pants.engine.fs import Digest, MergeDigests
 from pants.engine.internals.native_engine import Snapshot
 from pants.engine.internals.selectors import Get, MultiGet
@@ -17,6 +22,8 @@
 from pants.engine.unions import UnionRule
 from pants.option.option_types import SkipOption
 from pants.option.subsystem import Subsystem
+from pants.util.frozendict import FrozenDict
+from pants.util.strutil import pluralize
 
 logger = logging.getLogger(__name__)
 
@@ -34,14 +41,61 @@ class TffmtRequest(FmtTargetsRequest):
     name = TfFmtSubsystem.options_scope
 
 
+@dataclass(frozen=True)
+class TffmtPartitionRequest:
+    field_sets: tuple[TerraformFieldSet, ...]
+
+
+class TffmtPartitions(FrozenDict[Any, "tuple[str, ...]"]):
+    pass
+
+
+@dataclass(frozen=True)
+class TffmtSubPartitionRequest:
+    files: tuple[str, ...]
+    key: Any
+    snapshot: Snapshot  # NB: May contain more files than `files`
+
+
+@rule
+async def partition_tffmt(request: TffmtPartitionRequest, tffmt: TfFmtSubsystem) -> TffmtPartitions:
+    if tffmt.skip:
+        return TffmtPartitions()
+
+    source_files = await Get(
+        SourceFiles, SourceFilesRequest([field_set.sources for field_set in request.field_sets])
+    )
+
+    return TffmtPartitions(
+        (directory, tuple(files))
+        for directory, files in partition_files_by_directory(source_files.files).items()
+    )
+
+
+@rule
+async def run_tffmt(request: TffmtSubPartitionRequest) -> ProcessResult:
+    directory = cast(str, request.key)
+    snapshot = request.snapshot
+
+    result = await Get(
+        ProcessResult,
+        TerraformProcess(
+            args=("fmt", directory),
+            input_digest=snapshot.digest,
+            output_files=request.files,
+            description=f"Run `terraform fmt` on {pluralize(len(request.files), 'file')}.",
+        ),
+    )
+
+    return result
+
+
 @rule(desc="Format with `terraform fmt`")
 async def tffmt_fmt(request: TffmtRequest, tffmt: TfFmtSubsystem) -> FmtResult:
-    if tffmt.skip:
-        return FmtResult.skip(formatter_name=request.name)
-    setup = await Get(StyleSetup, StyleSetupRequest(request, ("fmt",)))
+    partitions = await Get(TffmtPartitions, TffmtPartitionRequest(request.field_sets))
     results = await MultiGet(
-        Get(ProcessResult, TerraformProcess, process)
-        for _, (process, _) in setup.directory_to_process.items()
+        Get(ProcessResult, TffmtSubPartitionRequest(files, key, request.snapshot))
+        for key, files in partitions.items()
     )
 
     def format(directory, output):
@@ -58,7 +112,7 @@ def format(directory, output):
 
     stdout_content = ""
     stderr_content = ""
-    for directory, result in zip(setup.directory_to_process.keys(), results):
+    for directory, result in zip(partitions, results):
         stdout_content += format(directory, result.stdout)
         stderr_content += format(directory, result.stderr)
 
@@ -67,7 +121,7 @@ def format(directory, output):
     output_snapshot = await Get(Snapshot, Digest, output_digest)
 
     fmt_result = FmtResult(
-        input=setup.original_snapshot,
+        input=request.snapshot,
         output=output_snapshot,
         stdout=stdout_content,
         stderr=stderr_content,

src/python/pants/backend/terraform/lint/tffmt/tffmt_integration_test.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from pants.backend.terraform import style, tool
+from pants.backend.terraform import tool
 from pants.backend.terraform.lint.tffmt import tffmt
 from pants.backend.terraform.lint.tffmt.tffmt import TffmtRequest
 from pants.backend.terraform.target_types import TerraformFieldSet, TerraformModuleTarget
@@ -52,7 +52,6 @@ def rule_runner() -> RuleRunner:
             *external_tool.rules(),
             *tffmt.rules(),
             *tool.rules(),
-            *style.rules(),
             *source_files.rules(),
             QueryRule(FmtResult, (TffmtRequest,)),
             QueryRule(SourceFiles, (SourceFilesRequest,)),

src/python/pants/backend/terraform/partition.py

@@ -0,0 +1,18 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+import os
+from collections import defaultdict
+from typing import Iterable
+
+
+def partition_files_by_directory(filepaths: Iterable[str]) -> dict[str, list[str]]:
+    """Maps files to directories, since `terraform` operates on a directory-by-directory basis."""
+    directories = defaultdict(list)
+    for filepath in filepaths:
+        directory = os.path.dirname(filepath) or "."
+        directories[directory].append(filepath)
+
+    return directories