Use address as stable key when batching field sets in lint/test (…

…Cherry-pick of pantsbuild#18725) (pantsbuild#18733)

Closes pantsbuild#18566

Using `str(x)` as the stable batching key ends up using `repr` when `x`
is a field set. After pantsbuild#18719 this no longer raises an exception, but
it's still an unintended behavior change introduced as part of the new
partitioning logic in 2.15.x.

Revert back to using `.address` as the partitioning key for field sets,
and add regression tests.

src/python/pants/core/goals/lint.py

@@ -427,12 +427,10 @@ async def lint(
     if not partitions_by_request_type:
         return Lint(exit_code=0)
 
-    def batch_by_size(
-        iterable: Iterable[_T], key: Callable[[_T], str] = lambda x: str(x)
-    ) -> Iterator[tuple[_T, ...]]:
+    def batch_by_size(iterable: Iterable[_T]) -> Iterator[tuple[_T, ...]]:
         batches = partition_sequentially(
             iterable,
-            key=key,
+            key=lambda x: str(x.address) if isinstance(x, FieldSet) else str(x),
             size_target=lint_subsystem.batch_size,
             size_max=4 * lint_subsystem.batch_size,
         )

src/python/pants/core/goals/lint_test.py

@@ -32,7 +32,7 @@
 from pants.engine.fs import PathGlobs, SpecsPaths, Workspace
 from pants.engine.internals.native_engine import EMPTY_SNAPSHOT, Snapshot
 from pants.engine.rules import QueryRule
-from pants.engine.target import FieldSet, FilteredTargets, MultipleSourcesField, Target
+from pants.engine.target import Field, FieldSet, FilteredTargets, MultipleSourcesField, Target
 from pants.engine.unions import UnionMembership
 from pants.option.option_types import SkipOption
 from pants.option.subsystem import Subsystem
@@ -48,15 +48,21 @@ class MockMultipleSourcesField(MultipleSourcesField):
     pass
 
 
+class MockRequiredField(Field):
+    alias = "required"
+    required = True
+
+
 class MockTarget(Target):
     alias = "mock_target"
-    core_fields = (MockMultipleSourcesField,)
+    core_fields = (MockMultipleSourcesField, MockRequiredField)
 
 
 @dataclass(frozen=True)
 class MockLinterFieldSet(FieldSet):
     required_fields = (MultipleSourcesField,)
     sources: MultipleSourcesField
+    required: MockRequiredField
 
 
 class MockLintRequest(LintRequest, metaclass=ABCMeta):
@@ -259,7 +265,9 @@ def rule_runner() -> RuleRunner:
 
 
 def make_target(address: Optional[Address] = None) -> Target:
-    return MockTarget({}, address or Address("", target_name="tests"))
+    return MockTarget(
+        {MockRequiredField.alias: "present"}, address or Address("", target_name="tests")
+    )
 
 
 def run_lint_rule(
@@ -511,8 +519,16 @@ class LintKitchenRequest(LintTargetsRequest):
     ]
     rule_runner = RuleRunner(rules=rules)
     field_sets = (
-        MockLinterFieldSet(Address("knife"), MultipleSourcesField(["knife"], Address("knife"))),
-        MockLinterFieldSet(Address("bowl"), MultipleSourcesField(["bowl"], Address("bowl"))),
+        MockLinterFieldSet(
+            Address("knife"),
+            MultipleSourcesField(["knife"], Address("knife")),
+            MockRequiredField("present", Address("")),
+        ),
+        MockLinterFieldSet(
+            Address("bowl"),
+            MultipleSourcesField(["bowl"], Address("bowl")),
+            MockRequiredField("present", Address("")),
+        ),
     )
     partitions = rule_runner.request(Partitions, [LintKitchenRequest.PartitionRequest(field_sets)])
     assert len(partitions) == 1

src/python/pants/core/goals/test.py

@@ -705,7 +705,7 @@ def partitions_get(request_type: type[TestRequest]) -> Get[Partitions]:
         for partition in partitions
         for batch in partition_sequentially(
             partition.elements,
-            key=lambda x: str(x),
+            key=lambda x: str(x.address) if isinstance(x, FieldSet) else str(x),
             size_target=test_subsystem.batch_size,
             size_max=2 * test_subsystem.batch_size,
         )

src/python/pants/core/goals/test_test.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from abc import ABCMeta, abstractmethod
+from abc import abstractmethod
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
@@ -60,6 +60,7 @@
 from pants.engine.process import InteractiveProcess, InteractiveProcessResult, ProcessResultMetadata
 from pants.engine.target import (
     BoolField,
+    Field,
     MultipleSourcesField,
     Target,
     TargetRootsToFieldSets,
@@ -93,9 +94,14 @@ class MockSkipTestsField(BoolField):
     default = False
 
 
+class MockRequiredField(Field):
+    alias = "required"
+    required = True
+
+
 class MockTarget(Target):
     alias = "mock_target"
-    core_fields = (MockMultipleSourcesField, MockSkipTestsField)
+    core_fields = (MockMultipleSourcesField, MockSkipTestsField, MockRequiredField)
 
 
 @dataclass(frozen=True)
@@ -107,8 +113,11 @@ class MockCoverageDataCollection(CoverageDataCollection):
     element_type = MockCoverageData
 
 
-class MockTestFieldSet(TestFieldSet, metaclass=ABCMeta):
-    required_fields = (MultipleSourcesField,)
+@dataclass(frozen=True)
+class MockTestFieldSet(TestFieldSet):
+    required_fields = (MultipleSourcesField, MockRequiredField)
+    sources: MultipleSourcesField
+    required: MockRequiredField
 
     @classmethod
     def opt_out(cls, tgt: Target) -> bool:
@@ -196,7 +205,7 @@ def rule_runner() -> RuleRunner:
 def make_target(address: Address | None = None, *, skip: bool = False) -> Target:
     if address is None:
         address = Address("", target_name="tests")
-    return MockTarget({MockSkipTestsField.alias: skip}, address)
+    return MockTarget({MockSkipTestsField.alias: skip, MockRequiredField.alias: "present"}, address)
 
 
 def run_test_rule(