Mypy and Pylint partition via CoarsenedTargets to allow graph computa…

…tion to be shared.

[ci skip-rust]
[ci skip-build-wheels]

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

@@ -3,19 +3,15 @@
 
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Tuple
+from typing import Mapping, Set, Tuple
 
 from pants.backend.python.lint.pylint.subsystem import (
     Pylint,
     PylintFieldSet,
     PylintFirstPartyPlugins,
 )
 from pants.backend.python.subsystems.setup import PythonSetup
-from pants.backend.python.target_types import (
-    InterpreterConstraintsField,
-    PythonResolveField,
-    PythonSourceField,
-)
+from pants.backend.python.target_types import InterpreterConstraintsField, PythonResolveField
 from pants.backend.python.util_rules import pex_from_targets
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.pex import (
@@ -37,7 +33,7 @@
 from pants.engine.fs import CreateDigest, Digest, Directory, MergeDigests, RemovePrefix
 from pants.engine.process import FallibleProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
-from pants.engine.target import Target, TransitiveTargets, TransitiveTargetsRequest
+from pants.engine.target import CoarsenedTarget, CoarsenedTargets, CoarsenedTargetsRequest, Target
 from pants.engine.unions import UnionRule
 from pants.util.logging import LogLevel
 from pants.util.ordered_set import FrozenOrderedSet, OrderedSet
@@ -46,7 +42,7 @@
 
 @dataclass(frozen=True)
 class PylintPartition:
-    root_targets: FrozenOrderedSet[Target]
+    root_field_sets: FrozenOrderedSet[PylintFieldSet]
     closure: FrozenOrderedSet[Target]
     interpreter_constraints: InterpreterConstraints
 
@@ -77,7 +73,7 @@ async def pylint_lint_partition(
     requirements_pex_get = Get(
         Pex,
         RequirementsPexRequest(
-            (t.address for t in partition.root_targets),
+            (fs.address for fs in partition.root_field_sets),
             # NB: These constraints must be identical to the other PEXes. Otherwise, we risk using
             # a different version for the requirements than the other two PEXes, which can result
             # in a PEX runtime error about missing dependencies.
@@ -96,7 +92,7 @@ async def pylint_lint_partition(
 
     prepare_python_sources_get = Get(PythonSourceFiles, PythonSourceFilesRequest(partition.closure))
     field_set_sources_get = Get(
-        SourceFiles, SourceFilesRequest(t[PythonSourceField] for t in partition.root_targets)
+        SourceFiles, SourceFilesRequest(fs.source for fs in partition.root_field_sets)
     )
     # Ensure that the empty report dir exists.
     report_directory_digest_get = Get(Digest, CreateDigest([Directory(REPORT_DIR)]))
@@ -161,8 +157,8 @@ async def pylint_lint_partition(
             input_digest=input_digest,
             output_directories=(REPORT_DIR,),
             extra_env={"PEX_EXTRA_SYS_PATH": ":".join(pythonpath)},
-            concurrency_available=len(partition.root_targets),
-            description=f"Run Pylint on {pluralize(len(partition.root_targets), 'file')}.",
+            concurrency_available=len(partition.root_field_sets),
+            description=f"Run Pylint on {pluralize(len(partition.root_field_sets), 'file')}.",
             level=LogLevel.DEBUG,
         ),
     )
@@ -188,48 +184,50 @@ async def pylint_determine_partitions(
     # Note that Pylint uses the AST of the interpreter that runs it. So, we include any plugin
     # targets in this interpreter constraints calculation. However, we don't have to consider the
     # resolve of the plugin targets, per https://github.com/pantsbuild/pants/issues/14320.
-    transitive_targets_per_field_set = await MultiGet(
-        Get(TransitiveTargets, TransitiveTargetsRequest([field_set.address]))
-        for field_set in request.field_sets
+    coarsened_targets = await Get(
+        CoarsenedTargets,
+        CoarsenedTargetsRequest(
+            (field_set.address for field_set in request.field_sets), expanded_targets=True
+        ),
     )
 
-    resolve_and_interpreter_constraints_to_transitive_targets = defaultdict(set)
-    for transitive_targets in transitive_targets_per_field_set:
-        resolve = transitive_targets.roots[0][PythonResolveField].normalized_value(python_setup)
+    resolve_and_interpreter_constraints_to_coarsened_targets: Mapping[
+        tuple[str, InterpreterConstraints],
+        tuple[OrderedSet[PylintFieldSet], OrderedSet[Target], Set[CoarsenedTarget]],
+    ] = defaultdict(lambda: (OrderedSet(), OrderedSet(), set()))
+    for root, ct in zip(request.field_sets, coarsened_targets):
+        # NB: If there is a cycle in the roots, we still only take the first resolve, as the other
+        # members will be validated when the partition is actually built.
+        resolve = ct.representative[PythonResolveField].normalized_value(python_setup)
         interpreter_constraints = InterpreterConstraints.create_from_compatibility_fields(
             (
                 *(
                     tgt[InterpreterConstraintsField]
-                    for tgt in transitive_targets.closure
+                    for tgt in ct.closure()
                     if tgt.has_field(InterpreterConstraintsField)
                 ),
                 *first_party_plugins.interpreter_constraints_fields,
             ),
             python_setup,
         )
-        resolve_and_interpreter_constraints_to_transitive_targets[
+        roots, closure, visited = resolve_and_interpreter_constraints_to_coarsened_targets[
             (resolve, interpreter_constraints)
-        ].add(transitive_targets)
-
-    partitions = []
-    for (_resolve, interpreter_constraints), all_transitive_targets in sorted(
-        resolve_and_interpreter_constraints_to_transitive_targets.items()
-    ):
-        combined_roots: OrderedSet[Target] = OrderedSet()
-        combined_closure: OrderedSet[Target] = OrderedSet()
-        for transitive_targets in all_transitive_targets:
-            combined_roots.update(transitive_targets.roots)
-            combined_closure.update(transitive_targets.closure)
-        partitions.append(
-            # Note that we don't need to pass the resolve. pex_from_targets.py will already
-            # calculate it by inspecting the roots & validating that all dependees are valid.
-            PylintPartition(
-                FrozenOrderedSet(combined_roots),
-                FrozenOrderedSet(combined_closure),
-                interpreter_constraints,
-            )
+        ]
+        roots.add(root)
+        closure.update(ct.closure(visited))
+
+    return PylintPartitions(
+        # Note that we don't need to pass the resolve. pex_from_targets.py will already
+        # calculate it by inspecting the roots & validating that all dependees are valid.
+        PylintPartition(
+            FrozenOrderedSet(roots),
+            FrozenOrderedSet(closure),
+            interpreter_constraints,
         )
-    return PylintPartitions(partitions)
+        for (_resolve, interpreter_constraints), (roots, closure, _) in sorted(
+            resolve_and_interpreter_constraints_to_coarsened_targets.items()
+        )
+    )
 
 
 @rule(desc="Lint using Pylint", level=LogLevel.DEBUG)

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

@@ -518,7 +518,7 @@ def assert_partition(
         partition: PylintPartition, roots: list[Target], deps: list[Target], interpreter: str
     ) -> None:
         root_addresses = {t.address for t in roots}
-        assert {t.address for t in partition.root_targets} == root_addresses
+        assert {fs.address for fs in partition.root_field_sets} == root_addresses
         assert {t.address for t in partition.closure} == {
             *root_addresses,
             *(t.address for t in deps),

src/python/pants/backend/python/typecheck/mypy/rules.py

@@ -4,7 +4,7 @@
 import itertools
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Iterable, Optional, Tuple
+from typing import Iterable, Mapping, Optional, Set, Tuple
 
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import PythonResolveField, PythonSourceField
@@ -29,7 +29,13 @@
 from pants.engine.fs import CreateDigest, Digest, FileContent, MergeDigests, RemovePrefix
 from pants.engine.process import FallibleProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
-from pants.engine.target import FieldSet, Target, TransitiveTargets, TransitiveTargetsRequest
+from pants.engine.target import (
+    CoarsenedTarget,
+    CoarsenedTargets,
+    CoarsenedTargetsRequest,
+    FieldSet,
+    Target,
+)
 from pants.engine.unions import UnionRule
 from pants.util.logging import LogLevel
 from pants.util.ordered_set import FrozenOrderedSet, OrderedSet
@@ -49,7 +55,7 @@ def opt_out(cls, tgt: Target) -> bool:
 
 @dataclass(frozen=True)
 class MyPyPartition:
-    root_targets: FrozenOrderedSet[Target]
+    root_field_sets: FrozenOrderedSet[MyPyFieldSet]
     closure: FrozenOrderedSet[Target]
     interpreter_constraints: InterpreterConstraints
 
@@ -125,14 +131,14 @@ async def mypy_typecheck_partition(
     closure_sources_get = Get(PythonSourceFiles, PythonSourceFilesRequest(partition.closure))
     roots_sources_get = Get(
         SourceFiles,
-        SourceFilesRequest(tgt.get(PythonSourceField) for tgt in partition.root_targets),
+        SourceFilesRequest(fs.sources for fs in partition.root_field_sets),
     )
 
     # See `requirements_venv_pex` for how this will get wrapped in a `VenvPex`.
     requirements_pex_get = Get(
         Pex,
         RequirementsPexRequest(
-            (tgt.address for tgt in partition.root_targets),
+            (fs.address for fs in partition.root_field_sets),
             hardcoded_interpreter_constraints=partition.interpreter_constraints,
         ),
     )
@@ -255,41 +261,43 @@ async def mypy_determine_partitions(
 ) -> MyPyPartitions:
     # When determining how to batch by interpreter constraints, we must consider the entire
     # transitive closure to get the final resulting constraints.
-    transitive_targets_per_field_set = await MultiGet(
-        Get(TransitiveTargets, TransitiveTargetsRequest([field_set.address]))
-        for field_set in request.field_sets
+    coarsened_targets = await Get(
+        CoarsenedTargets,
+        CoarsenedTargetsRequest(
+            (field_set.address for field_set in request.field_sets), expanded_targets=True
+        ),
     )
 
-    resolve_and_interpreter_constraints_to_transitive_targets = defaultdict(set)
-    for transitive_targets in transitive_targets_per_field_set:
-        resolve = transitive_targets.roots[0][PythonResolveField].normalized_value(python_setup)
+    resolve_and_interpreter_constraints_to_coarsened_targets: Mapping[
+        tuple[str, InterpreterConstraints],
+        tuple[OrderedSet[MyPyFieldSet], OrderedSet[Target], Set[CoarsenedTarget]],
+    ] = defaultdict(lambda: (OrderedSet(), OrderedSet(), set()))
+    for root, ct in zip(request.field_sets, coarsened_targets):
+        # NB: If there is a cycle in the roots, we still only take the first resolve, as the other
+        # members will be validated when the partition is actually built.
+        resolve = ct.representative[PythonResolveField].normalized_value(python_setup)
         interpreter_constraints = (
-            InterpreterConstraints.create_from_targets(transitive_targets.closure, python_setup)
+            InterpreterConstraints.create_from_targets(ct.closure(), python_setup)
             or mypy.interpreter_constraints
         )
-        resolve_and_interpreter_constraints_to_transitive_targets[
+        roots, closure, visited = resolve_and_interpreter_constraints_to_coarsened_targets[
             (resolve, interpreter_constraints)
-        ].add(transitive_targets)
-
-    partitions = []
-    for (_resolve, interpreter_constraints), all_transitive_targets in sorted(
-        resolve_and_interpreter_constraints_to_transitive_targets.items()
-    ):
-        combined_roots: OrderedSet[Target] = OrderedSet()
-        combined_closure: OrderedSet[Target] = OrderedSet()
-        for transitive_targets in all_transitive_targets:
-            combined_roots.update(transitive_targets.roots)
-            combined_closure.update(transitive_targets.closure)
-        partitions.append(
-            # Note that we don't need to pass the resolve. pex_from_targets.py will already
-            # calculate it by inspecting the roots & validating that all dependees are valid.
-            MyPyPartition(
-                FrozenOrderedSet(combined_roots),
-                FrozenOrderedSet(combined_closure),
-                interpreter_constraints,
-            )
+        ]
+        roots.add(root)
+        closure.update(ct.closure(visited))
+
+    return MyPyPartitions(
+        # Note that we don't need to pass the resolve. pex_from_targets.py will already
+        # calculate it by inspecting the roots & validating that all dependees are valid.
+        MyPyPartition(
+            FrozenOrderedSet(roots),
+            FrozenOrderedSet(closure),
+            interpreter_constraints,
         )
-    return MyPyPartitions(partitions)
+        for (_resolve, interpreter_constraints), (roots, closure, _) in sorted(
+            resolve_and_interpreter_constraints_to_coarsened_targets.items()
+        )
+    )
 
 
 # TODO(#10864): Improve performance, e.g. by leveraging the MyPy cache.

src/python/pants/backend/python/typecheck/mypy/rules_integration_test.py

@@ -756,7 +756,7 @@ def assert_partition(
         partition: MyPyPartition, roots: list[Target], deps: list[Target], interpreter: str
     ) -> None:
         root_addresses = {t.address for t in roots}
-        assert {t.address for t in partition.root_targets} == root_addresses
+        assert {fs.address for fs in partition.root_field_sets} == root_addresses
         assert {t.address for t in partition.closure} == {
             *root_addresses,
             *(t.address for t in deps),

src/python/pants/engine/internals/graph.py

@@ -552,7 +552,9 @@ async def coarsened_targets(request: CoarsenedTargetsRequest) -> CoarsenedTarget
     dependency_mapping = await Get(
         _DependencyMapping,
         _DependencyMappingRequest(
-            TransitiveTargetsRequest(request.roots, include_special_cased_deps=request.include_special_cased_deps),
+            TransitiveTargetsRequest(
+                request.roots, include_special_cased_deps=request.include_special_cased_deps
+            ),
             expanded_targets=request.expanded_targets,
         ),
     )

src/python/pants/engine/target.py

@@ -759,7 +759,11 @@ class CoarsenedTargetsRequest:
     include_special_cased_deps: bool
 
     def __init__(
-            self, roots: Iterable[Address], *, expanded_targets: bool = False, include_special_cased_deps: bool = False
+        self,
+        roots: Iterable[Address],
+        *,
+        expanded_targets: bool = False,
+        include_special_cased_deps: bool = False,
     ) -> None:
         self.roots = tuple(roots)
         self.expanded_targets = expanded_targets