nfpm: support chaining multiple InjectNfpmPackageFieldsRequests (sorted w/ simple priority integer) (#22864)

This extends then `pants.backend.experimental.nfpm` backend's plugin
API, allowing multiple implementations of this polymorphic-rule plugin
hook (before this change, only one implementation was allowed):

- `inject_nfpm_package_fields(InjectNfpmPackageFieldsRequest) -> InjectedNfpmPackageFields`

Unlike many polymorphic/union rules, each rule runs sequentially instead
of concurrently. This is conceptually similar to middleware, because the
request passed to each rule includes the results of the previous (lower
priority) rule. So, one rule can override or extend fields that were
injected in previous rules, not just the original fields from the
BUILD-file-based `nfpm_*_package` target.

I tried several approaches to sorting the request classes, but most were
error prone or simply didn't work. Heuristic-based sorting proved to be
complex and didn't work in all of my test cases (eg one heuristic was
looking at the class's module to make pants-provided requests lower
priority). A simple `priority` integer is far simpler and more reliable,
so that's what this implements.

Another thing that didn't work was adding `__lt__` as a `@classmethod`
on the abstract `InjectNfpmPackageFieldsRequest`. That fails because
`sorted()` requires an instance method, not a class method. The only way
I found to add an `__lt__` method on a class is via a metaclass. So,
this adds a `_PrioritizedSortableClassMetaclass` to
`InjectNfpmPackageFieldsRequest` allowing for simple `OneRequest <
TwoRequest` sorting of the class type like this:

```diff
     inject_nfpm_config_request_types = union_membership.get(InjectNfpmPackageFieldsRequest)
     applicable_inject_nfpm_config_request_types = tuple(
-        request
-        for request in inject_nfpm_config_request_types
-        if request.is_applicable(target)
+        sorted(
+            request_type
+            for request_type in inject_nfpm_config_request_types
+            if request_type.is_applicable(target)
+        )
     )
```

Then, to make subclasses higher priority than their parent classes, I
added an `__init_subclass__` method that increases `cls.priority`
(unless the subclass already has a different priority than the parent
class). It's a little odd to use both a metaclass and
`__init_subclass__`, but I did not like the idea of re-implementing the
on-subclass-creation logic in our metaclass.

`InjectNfpmPackageFieldsRequest` was initially modeled on the python
backend's `SetupKwargsRequest` [1].
If the rule-chaining proves itself in the experimental `nfpm` backend,
maybe we can move `_PrioritizedSortableClassMetaclass` out of the
backend to somewhere in pants core, updating `SetupKwargsRequest` and
related rules to make use of rule-chaining as well.

[1] https://www.pantsbuild.org/stable/docs/writing-plugins/common-plugin-tasks/custom-python-artifact-kwargs

Author: cognifloyd

docs/notes/2.31.x.md

@@ -46,6 +46,8 @@ Pants no longer supports loading `pkg_resources`-style namespace packages for pl
 
 Added a new rule to help in-repo plugins implement the `inject_nfpm_package_fields(InjectNfpmPackageFieldsRequest) -> InjectedNfpmPackageFields` polymorphic rule. The `get_package_field_sets_for_nfpm_content_file_deps` rule (in the `pants.backend.nfpm.util_rules.contents` module) collects selected `PackageFieldSet`s from the contents of an `nfpm_*_package` so that the packages can be analyzed to inject things like package requirements.
 
+Added prioritized rule chaining for `InjectNfpmPackageFieldsRequest` polymorphic rule implementations. Before, only one request type was allowed to inject fields on each `nfpm_*_package` target. Now, the `InjectNfpmPackageFieldsRequest` subclasses can declare a priority (as `ClassVar[int]`) so that the rule runs after lower-priority requests and before higher-priority requests. By default, `InjectNfpmPackageFieldsRequest` implementations are priority 10, and any subclasses have that priority increased by 1 so that the subclass is higher priority than its parent class. Any implementations distributed with pants should have a lower priority (`< 10`) so that in-repo and external plugins can override the pants-provided defaults.
+
 ## Full Changelog
 
 For the full changelog, see the individual GitHub Releases for this series: <https://github.com/pantsbuild/pants/releases>

src/python/pants/backend/nfpm/util_rules/inject_config.py

@@ -3,9 +3,10 @@
 
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
+from abc import ABC, ABCMeta, abstractmethod
 from collections.abc import Iterable
 from dataclasses import dataclass
+from typing import Any, ClassVar, TypeVar, cast
 
 from pants.backend.nfpm.fields.scripts import NfpmPackageScriptsField
 from pants.engine.environment import EnvironmentName
@@ -16,6 +17,9 @@
 from pants.util.frozendict import FrozenDict
 from pants.util.strutil import softwrap
 
+# NB: This TypeVar serves the same purpose here as in pants.engine.target
+_F = TypeVar("_F", bound=Field)
+
 
 @dataclass(frozen=True)
 class InjectedNfpmPackageFields:
@@ -71,27 +75,87 @@ def __init__(
         )
 
 
+class _PrioritizedSortableClassMetaclass(ABCMeta):
+    """This metaclass implements prioritized sorting of subclasses (not class instances)."""
+
+    priority: ClassVar[int]
+
+    def __lt__(self, other: Any) -> bool:
+        """Determine if this class is lower priority than `other` (when chaining request rules).
+
+        The rule that runs the lowest priority request goes first, and the rule that runs the
+        highest priority request goes last. The results (the `injected_fields`) of lower priority
+        rules can be overridden by higher priority rules. The last rule to run, the rule for the
+        highest priority request class, can override any of the fields injected by lower priority
+        request rules.
+        """
+        if not isinstance(other, _PrioritizedSortableClassMetaclass):
+            return NotImplemented
+        if self.priority != other.priority:
+            return self.priority < other.priority
+        # other has same priority: fall back to name comparison (ensures deterministic sort)
+        return (self.__module__, self.__qualname__) < (other.__module__, other.__qualname__)
+
+
 # Note: This only exists as a hook for additional logic for nFPM config generation, e.g. for plugin
 # authors. To resolve `InjectedNfpmPackageFields`, call `determine_injected_nfpm_package_fields`,
 # which handles running any custom implementations vs. using the default implementation.
 @union(in_scope_types=[EnvironmentName])
 @dataclass(frozen=True)
-class InjectNfpmPackageFieldsRequest(ABC):
+class InjectNfpmPackageFieldsRequest(ABC, metaclass=_PrioritizedSortableClassMetaclass):
     """A request to inject nFPM config for nfpm_package_* targets.
 
     By default, Pants will use the nfpm_package_* fields in the BUILD file unchanged to generate the
     nfpm.yaml config file for nFPM. To customize this, subclass `InjectNfpmPackageFieldsRequest`,
     register `UnionRule(InjectNfpmPackageFieldsRequest, MyCustomInjectNfpmPackageFieldsRequest)`,
     and add a rule that takes your subclass as a parameter and returns `InjectedNfpmPackageFields`.
+
+    The `target` attribute of this class holds the original target as defined in BUILD files.
+    The `injected_fields` attribute of this class contains the results of any previous rule.
+    `injected_fields` will be empty for the first rule in the chain. Subsequent rules can remove
+    or replace fields injected by previous rules. The final rule in the chain returns the final
+    `InjectedNfpmPackageFields` instance that is used to actually generate the nfpm config.
+    In general, rules should include a copy of `request.injected_fields` in their return value
+    with something like this:
+
+        address = request.target.address
+        fields: list[Field] = list(request.injected_fields.values())
+        fields.append(NfpmFoobarField("foobar", address))
+        return InjectedNfpmPackageFields(fields, address=address)
+
+    Chaining rules like this allows pants to inject some fields, while allowing in-repo plugins
+    to override or remove them.
     """
 
     target: Target
+    injected_fields: FrozenDict[type[Field], Field]
+
+    # Classes in pants-provided backends should be priority<10 so that in-repo and external
+    # plugins are higher priority by default. This way, the fields that get injected by default
+    # can be overridden as needed in the in-repo or external plugins.
+    priority: ClassVar[int] = 10
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        if cls.priority == getattr(super(), "priority", cls.priority):
+            # subclasses are higher priority than their parent class (unless set otherwise).
+            cls.priority += 1
 
     @classmethod
     @abstractmethod
     def is_applicable(cls, target: Target) -> bool:
         """Whether to use this InjectNfpmPackageFieldsRequest implementation for this target."""
 
+    def get_field(self, field: type[_F]) -> _F:
+        """Get a `Field` from `injected_fields` (returned by earlier rules) or from the target.
+
+        This will throw a KeyError if the `Field` is not registered on the target (unless an earlier
+        rule added it to `injected_fields` which might be disallowed in the future).
+        """
+        if field in self.injected_fields:
+            return cast(_F, self.injected_fields[field])
+        return self.target[field]
+
 
 @rule(polymorphic=True)
 async def inject_nfpm_package_fields(
@@ -115,38 +179,33 @@ async def determine_injected_nfpm_package_fields(
     wrapper: NfpmPackageTargetWrapper, union_membership: UnionMembership
 ) -> InjectedNfpmPackageFields:
     target = wrapper.target
-    inject_nfpm_config_requests = union_membership.get(InjectNfpmPackageFieldsRequest)
-    applicable_inject_nfpm_config_requests = tuple(
-        request for request in inject_nfpm_config_requests if request.is_applicable(target)
+    inject_nfpm_config_request_types = union_membership.get(InjectNfpmPackageFieldsRequest)
+
+    # Requests are sorted (w/ priority ClassVar) before chaining the rules that take them.
+    applicable_inject_nfpm_config_request_types = tuple(
+        sorted(
+            request_type
+            for request_type in inject_nfpm_config_request_types
+            if request_type.is_applicable(target)
+        )
     )
 
     # If no provided implementations, fall back to our default implementation that simply returns
     # what the user explicitly specified in the BUILD file.
-    if not applicable_inject_nfpm_config_requests:
+    if not applicable_inject_nfpm_config_request_types:
         return InjectedNfpmPackageFields((), address=target.address)
 
-    if len(applicable_inject_nfpm_config_requests) > 1:
-        possible_requests = sorted(
-            plugin.__name__ for plugin in applicable_inject_nfpm_config_requests
+    injected: InjectedNfpmPackageFields
+    injected_fields: FrozenDict[type[Field], Field] = FrozenDict()
+    for request_type in applicable_inject_nfpm_config_request_types:
+        chained_request: InjectNfpmPackageFieldsRequest = request_type(target, injected_fields)  # type: ignore[abstract]
+        injected = await inject_nfpm_package_fields(
+            **implicitly({chained_request: InjectNfpmPackageFieldsRequest})
         )
-        raise ValueError(
-            softwrap(
-                f"""
-                Multiple registered `InjectNfpmPackageFieldsRequest`s can work on the target
-                {target.address}, and it's ambiguous which to use: {possible_requests}
-
-                Please activate fewer implementations, or make the classmethod `is_applicable()`
-                more precise so that only one implementation is applicable for this target.
-                """
-            )
-        )
-    inject_nfpm_config_request_type = applicable_inject_nfpm_config_requests[0]
-    inject_nfpm_config_request: InjectNfpmPackageFieldsRequest = inject_nfpm_config_request_type(
-        target
-    )  # type: ignore[abstract]
-    return await inject_nfpm_package_fields(
-        **implicitly({inject_nfpm_config_request: InjectNfpmPackageFieldsRequest})
-    )
+        injected_fields = injected.field_values
+
+    # return the last result in the chain of results
+    return injected
 
 
 def rules():

src/python/pants/backend/nfpm/util_rules/inject_config_test.py

@@ -4,9 +4,11 @@
 from __future__ import annotations
 
 from textwrap import dedent
+from typing import Any
 
 import pytest
 
+from pants.backend.nfpm.fields.all import NfpmPlatformField
 from pants.backend.nfpm.fields.version import NfpmVersionField, NfpmVersionReleaseField
 from pants.backend.nfpm.target_types import target_types as nfpm_target_types
 from pants.backend.nfpm.target_types_rules import rules as nfpm_target_types_rules
@@ -24,24 +26,104 @@
 _PKG_NAME = "pkg"
 
 
-class PluginInjectNfpmPackageFieldsRequest(InjectNfpmPackageFieldsRequest):
+class PluginInjectFieldsRequest(InjectNfpmPackageFieldsRequest):
     @classmethod
     def is_applicable(cls, _) -> bool:
         return True
 
 
 @rule
 async def inject_nfpm_package_fields_plugin(
-    request: PluginInjectNfpmPackageFieldsRequest,
+    request: PluginInjectFieldsRequest,
 ) -> InjectedNfpmPackageFields:
     address = request.target.address
-    fields: list[Field] = [
-        NfpmVersionField("9.8.7-dev+git", address),
-        NfpmVersionReleaseField(6, address),
-    ]
+    # preserve fields from earlier rules in chain
+    fields: list[Field] = list(request.injected_fields.values())
+    if NfpmVersionField not in request.injected_fields:
+        fields.extend(
+            [
+                NfpmVersionField("9.8.7-dev+git", address),
+                NfpmVersionReleaseField(6, address),
+            ]
+        )
     return InjectedNfpmPackageFields(fields, address=address)
 
 
+class SubclassPluginInjectFieldsRequest(PluginInjectFieldsRequest):
+    pass
+
+
+class AnotherSubclassPluginInjectFieldsRequest(PluginInjectFieldsRequest):
+    pass
+
+
+@rule
+async def inject_nfpm_package_fields_subclass(
+    request: SubclassPluginInjectFieldsRequest,
+) -> InjectedNfpmPackageFields:
+    address = request.target.address
+    # preserve fields from earlier rules in chain
+    fields: list[Field] = list(request.injected_fields.values())
+    if not fields or NfpmVersionReleaseField in request.injected_fields:
+        release = 0
+        if NfpmVersionReleaseField in request.injected_fields:
+            old_release = request.injected_fields[NfpmVersionReleaseField].value
+            assert old_release is not None
+            release = 10 + old_release
+        fields.append(NfpmVersionReleaseField(release, address))
+    return InjectedNfpmPackageFields(fields, address=address)
+
+
+class HighPriorityInjectFieldsRequest(InjectNfpmPackageFieldsRequest):
+    priority = 100
+
+    @classmethod
+    def is_applicable(cls, _) -> bool:
+        return True
+
+
+@rule
+async def inject_nfpm_package_fields_high_priority(
+    request: HighPriorityInjectFieldsRequest,
+) -> InjectedNfpmPackageFields:
+    address = request.target.address
+    # preserve fields from earlier rules in chain
+    fields: list[Field] = list(request.injected_fields.values())
+    if not fields or NfpmVersionField not in request.injected_fields:
+        fields.extend(
+            [
+                NfpmVersionField("9.9.9-dev+git", address),
+                NfpmVersionReleaseField(9, address),
+            ]
+        )
+    # This high priority implementation wants to force Platform to always be "foobar"
+    # even if a lower priority rule injected NfpmPlatformField with a different value.
+    fields.append(NfpmPlatformField("foobar", address))
+    return InjectedNfpmPackageFields(fields, address=address)
+
+
+@pytest.mark.parametrize(
+    "a,b,expected_lt",
+    (
+        # HighPriority* > SubclassPlugin* > Plugin*
+        (HighPriorityInjectFieldsRequest, SubclassPluginInjectFieldsRequest, False),
+        (HighPriorityInjectFieldsRequest, PluginInjectFieldsRequest, False),
+        (SubclassPluginInjectFieldsRequest, HighPriorityInjectFieldsRequest, True),
+        (SubclassPluginInjectFieldsRequest, PluginInjectFieldsRequest, False),
+        (PluginInjectFieldsRequest, HighPriorityInjectFieldsRequest, True),
+        (PluginInjectFieldsRequest, SubclassPluginInjectFieldsRequest, True),
+        # Subclass* > AnotherSubclass* : with same priority, sort uses name of module+class
+        (SubclassPluginInjectFieldsRequest, AnotherSubclassPluginInjectFieldsRequest, False),
+        (AnotherSubclassPluginInjectFieldsRequest, SubclassPluginInjectFieldsRequest, True),
+        # self is equal which is not less than
+        (HighPriorityInjectFieldsRequest, HighPriorityInjectFieldsRequest, False),
+    ),
+)
+def test_nfpm_package_fields_request_class_sort(a: Any, b: Any, expected_lt):
+    ret = a < b
+    assert ret == expected_lt, f"a.priority={a.priority} b.priority={b.priority}"
+
+
 @pytest.fixture
 def rule_runner() -> RuleRunner:
     rule_runner = RuleRunner(
@@ -52,7 +134,11 @@ def rule_runner() -> RuleRunner:
             *nfpm_target_types_rules(),
             *nfpm_inject_config_rules(),
             inject_nfpm_package_fields_plugin,
-            UnionRule(InjectNfpmPackageFieldsRequest, PluginInjectNfpmPackageFieldsRequest),
+            inject_nfpm_package_fields_subclass,
+            inject_nfpm_package_fields_high_priority,
+            UnionRule(InjectNfpmPackageFieldsRequest, PluginInjectFieldsRequest),
+            UnionRule(InjectNfpmPackageFieldsRequest, SubclassPluginInjectFieldsRequest),
+            UnionRule(InjectNfpmPackageFieldsRequest, HighPriorityInjectFieldsRequest),
             QueryRule(InjectedNfpmPackageFields, (NfpmPackageTargetWrapper,)),
         ],
     )
@@ -90,6 +176,7 @@ def test_determine_injected_nfpm_package_fields(rule_runner: RuleRunner, package
     target = rule_runner.get_target(Address("", target_name=_PKG_NAME))
     result = rule_runner.request(InjectedNfpmPackageFields, [NfpmPackageTargetWrapper(target)])
     field_values = result.field_values
-    assert len(field_values) == 2
-    assert field_values[NfpmVersionField].value == "9.8.7-dev+git"
-    assert field_values[NfpmVersionReleaseField].value == 6
+    assert len(field_values) == 3
+    assert field_values[NfpmVersionField].value == "9.8.7-dev+git"  # (Plugin*)
+    assert field_values[NfpmVersionReleaseField].value == 16  # 6 (Plugin*) + 10 (SubclassPlugin*)
+    assert field_values[NfpmPlatformField].value == "foobar"  # (HighPriority*)