Fix performance for generating lockfiles for pytest and setuptools

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

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import itertools
 from dataclasses import dataclass
 
 from pants.backend.python.goals import lockfile
@@ -20,11 +19,11 @@
     InterpreterConstraintsField,
     PythonSourceField,
 )
-from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.backend.python.util_rules.partition import _find_all_unique_interpreter_constraints
 from pants.core.goals.generate_lockfiles import GenerateToolLockfileSentinel
 from pants.core.util_rules.config_files import ConfigFilesRequest
-from pants.engine.rules import Get, collect_rules, rule, rule_helper
-from pants.engine.target import AllTargets, AllTargetsRequest, FieldSet, Target
+from pants.engine.rules import collect_rules, rule
+from pants.engine.target import FieldSet, Target
 from pants.engine.unions import UnionRule
 from pants.option.option_types import ArgsListOption, FileOption, SkipOption
 from pants.util.docutil import git_url
@@ -84,25 +83,6 @@ def config_request(self) -> ConfigFilesRequest:
         )
 
 
-@rule_helper
-async def _bandit_interpreter_constraints(python_setup: PythonSetup) -> InterpreterConstraints:
-    # While Bandit will run in partitions, we need a set of constraints that works with every
-    # partition.
-    #
-    # This ORs all unique interpreter constraints. The net effect is that every possible Python
-    # interpreter used will be covered.
-    all_tgts = await Get(AllTargets, AllTargetsRequest())
-    unique_constraints = {
-        InterpreterConstraints.create_from_targets([tgt], python_setup)
-        for tgt in all_tgts
-        if BanditFieldSet.is_applicable(tgt)
-    }
-    constraints = InterpreterConstraints(
-        itertools.chain.from_iterable(ic for ic in unique_constraints if ic)
-    )
-    return constraints or InterpreterConstraints(python_setup.interpreter_constraints)
-
-
 class BanditLockfileSentinel(GeneratePythonToolLockfileSentinel):
     resolve_name = Bandit.options_scope
 
@@ -124,7 +104,7 @@ async def setup_bandit_lockfile(
             bandit, use_pex=python_setup.generate_lockfiles_with_pex
         )
 
-    constraints = await _bandit_interpreter_constraints(python_setup)
+    constraints = await _find_all_unique_interpreter_constraints(python_setup, BanditFieldSet)
     return GeneratePythonLockfile.from_tool(
         bandit,
         constraints,
@@ -150,7 +130,7 @@ async def bandit_export(
 ) -> ExportPythonTool:
     if not bandit.export:
         return ExportPythonTool(resolve_name=bandit.options_scope, pex_request=None)
-    constraints = await _bandit_interpreter_constraints(python_setup)
+    constraints = await _find_all_unique_interpreter_constraints(python_setup, BanditFieldSet)
     return ExportPythonTool(
         resolve_name=bandit.options_scope,
         pex_request=bandit.to_pex_request(interpreter_constraints=constraints),

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

@@ -1,12 +1,8 @@
 # Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-from dataclasses import dataclass
-
-from pants.backend.python.lint.black.skip_field import SkipBlackField
-from pants.backend.python.lint.black.subsystem import Black
+from pants.backend.python.lint.black.subsystem import Black, BlackFieldSet
 from pants.backend.python.subsystems.setup import PythonSetup
-from pants.backend.python.target_types import InterpreterConstraintsField, PythonSourceField
 from pants.backend.python.util_rules import pex
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.pex import PexRequest, VenvPex, VenvPexProcess
@@ -16,24 +12,11 @@
 from pants.engine.internals.native_engine import Snapshot
 from pants.engine.process import 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, softwrap
 
 
-@dataclass(frozen=True)
-class BlackFieldSet(FieldSet):
-    required_fields = (PythonSourceField,)
-
-    source: PythonSourceField
-    interpreter_constraints: InterpreterConstraintsField
-
-    @classmethod
-    def opt_out(cls, tgt: Target) -> bool:
-        return tgt.get(SkipBlackField).value
-
-
 class BlackRequest(FmtTargetsRequest):
     field_set_type = BlackFieldSet
     name = Black.options_scope

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

@@ -8,9 +8,9 @@
 import pytest
 
 from pants.backend.python import target_types_rules
-from pants.backend.python.lint.black.rules import BlackFieldSet, BlackRequest
+from pants.backend.python.lint.black.rules import BlackRequest
 from pants.backend.python.lint.black.rules import rules as black_rules
-from pants.backend.python.lint.black.subsystem import Black
+from pants.backend.python.lint.black.subsystem import Black, BlackFieldSet
 from pants.backend.python.lint.black.subsystem import rules as black_subsystem_rules
 from pants.backend.python.target_types import PythonSourcesGeneratorTarget
 from pants.core.goals.fmt import FmtResult

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

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import os.path
+from dataclasses import dataclass
 from typing import Iterable
 
 from pants.backend.python.goals import lockfile
@@ -15,19 +16,36 @@
 from pants.backend.python.lint.black.skip_field import SkipBlackField
 from pants.backend.python.subsystems.python_tool_base import ExportToolOption, PythonToolBase
 from pants.backend.python.subsystems.setup import PythonSetup
-from pants.backend.python.target_types import ConsoleScript
+from pants.backend.python.target_types import (
+    ConsoleScript,
+    InterpreterConstraintsField,
+    PythonSourceField,
+)
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.backend.python.util_rules.partition import _find_all_unique_interpreter_constraints
 from pants.core.goals.generate_lockfiles import GenerateToolLockfileSentinel
 from pants.core.util_rules.config_files import ConfigFilesRequest
-from pants.engine.rules import Get, collect_rules, rule, rule_helper
-from pants.engine.target import AllTargets, AllTargetsRequest
+from pants.engine.rules import collect_rules, rule, rule_helper
+from pants.engine.target import FieldSet, Target
 from pants.engine.unions import UnionRule
 from pants.option.option_types import ArgsListOption, BoolOption, FileOption, SkipOption
 from pants.util.docutil import git_url
 from pants.util.logging import LogLevel
 from pants.util.strutil import softwrap
 
 
+@dataclass(frozen=True)
+class BlackFieldSet(FieldSet):
+    required_fields = (PythonSourceField,)
+
+    source: PythonSourceField
+    interpreter_constraints: InterpreterConstraintsField
+
+    @classmethod
+    def opt_out(cls, tgt: Target) -> bool:
+        return tgt.get(SkipBlackField).value
+
+
 class Black(PythonToolBase):
     options_scope = "black"
     name = "Black"
@@ -92,10 +110,8 @@ async def _black_interpreter_constraints(
 ) -> InterpreterConstraints:
     constraints = black.interpreter_constraints
     if black.options.is_default("interpreter_constraints"):
-        all_tgts = await Get(AllTargets, AllTargetsRequest())
-        # TODO: fix to use `FieldSet.is_applicable()`.
-        code_constraints = InterpreterConstraints.create_from_targets(
-            (tgt for tgt in all_tgts if not tgt.get(SkipBlackField).value), python_setup
+        code_constraints = await _find_all_unique_interpreter_constraints(
+            python_setup, BlackFieldSet
         )
         if code_constraints is not None and code_constraints.requires_python38_or_newer(
             python_setup.interpreter_versions_universe

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

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import itertools
 from dataclasses import dataclass
 
 from pants.backend.python.goals import lockfile
@@ -22,7 +21,7 @@
     PythonSourceField,
 )
 from pants.backend.python.util_rules import python_sources
-from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.backend.python.util_rules.partition import _find_all_unique_interpreter_constraints
 from pants.backend.python.util_rules.pex_requirements import PexRequirements
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFilesRequest,
@@ -33,15 +32,8 @@
 from pants.engine.addresses import Addresses, UnparsedAddressInputs
 from pants.engine.fs import AddPrefix, Digest
 from pants.engine.internals.native_engine import EMPTY_DIGEST
-from pants.engine.rules import Get, collect_rules, rule, rule_helper
-from pants.engine.target import (
-    AllTargets,
-    AllTargetsRequest,
-    FieldSet,
-    Target,
-    TransitiveTargets,
-    TransitiveTargetsRequest,
-)
+from pants.engine.rules import Get, collect_rules, rule
+from pants.engine.target import FieldSet, Target, TransitiveTargets, TransitiveTargetsRequest
 from pants.engine.unions import UnionRule
 from pants.option.option_types import (
     ArgsListOption,
@@ -237,44 +229,6 @@ async def flake8_first_party_plugins(flake8: Flake8) -> Flake8FirstPartyPlugins:
     )
 
 
-# --------------------------------------------------------------------------------------
-# Interpreter constraints
-# --------------------------------------------------------------------------------------
-
-
-@rule_helper
-async def _flake8_interpreter_constraints(
-    first_party_plugins: Flake8FirstPartyPlugins,
-    python_setup: PythonSetup,
-) -> InterpreterConstraints:
-    # While Flake8 will run in partitions, we need a set of constraints that works with every
-    # partition.
-    #
-    # This ORs all unique interpreter constraints. The net effect is that every possible Python
-    # interpreter used will be covered.
-    all_tgts = await Get(AllTargets, AllTargetsRequest())
-    unique_constraints = {
-        InterpreterConstraints.create_from_compatibility_fields(
-            (
-                tgt[InterpreterConstraintsField],
-                *first_party_plugins.interpreter_constraints_fields,
-            ),
-            python_setup,
-        )
-        for tgt in all_tgts
-        if Flake8FieldSet.is_applicable(tgt)
-    }
-    if not unique_constraints:
-        unique_constraints.add(
-            InterpreterConstraints.create_from_compatibility_fields(
-                first_party_plugins.interpreter_constraints_fields,
-                python_setup,
-            )
-        )
-    constraints = InterpreterConstraints(itertools.chain.from_iterable(unique_constraints))
-    return constraints or InterpreterConstraints(python_setup.interpreter_constraints)
-
-
 # --------------------------------------------------------------------------------------
 # Lockfile
 # --------------------------------------------------------------------------------------
@@ -304,7 +258,11 @@ async def setup_flake8_lockfile(
             flake8, use_pex=python_setup.generate_lockfiles_with_pex
         )
 
-    constraints = await _flake8_interpreter_constraints(first_party_plugins, python_setup)
+    constraints = await _find_all_unique_interpreter_constraints(
+        python_setup,
+        Flake8FieldSet,
+        extra_constraints_per_tgt=first_party_plugins.interpreter_constraints_fields,
+    )
     return GeneratePythonLockfile.from_tool(
         flake8,
         constraints,
@@ -339,7 +297,11 @@ async def flake8_export(
 ) -> ExportPythonTool:
     if not flake8.export:
         return ExportPythonTool(resolve_name=flake8.options_scope, pex_request=None)
-    constraints = await _flake8_interpreter_constraints(first_party_plugins, python_setup)
+    constraints = await _find_all_unique_interpreter_constraints(
+        python_setup,
+        Flake8FieldSet,
+        extra_constraints_per_tgt=first_party_plugins.interpreter_constraints_fields,
+    )
     return ExportPythonTool(
         resolve_name=flake8.options_scope,
         pex_request=flake8.to_pex_request(

src/python/pants/backend/python/lint/flake8/subsystem_test.py

@@ -105,7 +105,7 @@ def test_first_party_plugins(rule_runner: RuleRunner) -> None:
     )
 
 
-def test_setup_lockfile_interpreter_constraints(rule_runner) -> None:
+def test_setup_lockfile(rule_runner) -> None:
     global_constraint = "==3.9.*"
 
     def assert_lockfile_request(