Allow plugin authors to attach custom fields to TargetData for pants peek

docs/docs/writing-plugins/the-target-api/concepts.mdx

@@ -196,3 +196,41 @@ class SyntheticExampleAllTargetsAtOnceRequest(SyntheticTargetsRequest):
 Any other default value for `path` should be considered invalid and yield undefined behaviour. (that is it may change without notice in future versions of Pants.)
 
 During rule execution, the `path` field of the `request` instance will hold the value for the path currently being parsed in case of a per directory mode of operation otherwise it will be `SyntheticTargetsRequest.SINGLE_REQUEST_FOR_ALL_TARGETS`.
+
+## Adding information to pants peek output
+
+Sometimes you may have metadata for a target that cannot be encompassed by a field, e.g. if it depends on the content of a file or if it requires some rule resolution to be calculated.
+
+You can attach this data to the output of `pants peek` by subclassing the `HasAdditionalTargetDataFieldSet` union type and register it as a union member with `UnionRule(HasAdditionalTargetDataFieldSet, SubclassedType)`. Then, implement a rule that takes `SubclassedType` as input and returns an `AdditionalTargetData` object.
+
+```python
+from dataclasses import dataclass
+from pants.backend.project_info.peek import AdditionalTargetData, HasAdditionalTargetDataFieldSet
+from pants.engine.unions import UnionRule
+from pants.engine.rules import collect_rules, rule
+
+
+@dataclass(frozen=True)
+class MyCustomTargetFieldSet(HasAdditionalTargetDataFieldSet):
+    ...
+
+
+@rule
+async def attach_custom_target_data(field_set: MyCustomTargetFieldSet) -> AdditionalTargetData:
+    # You can return any json-serializable type for the second field
+    return AdditionalTargetData("my_custom_target_data", {"hello": "world"})
+
+
+def rules():
+    return (*collect_rules(), UnionRule(HasAdditionalTargetDataFieldSet, MyCustomTargetFieldSet))
+```
+
+Then, if you run `pants peek --include-additional-info my/custom:target` you will see an `additional_info` field which will contain the following JSON object:
+
+```json
+{
+    "my_custom_target_data": {
+        "hello": "world"
+    }
+}
+```

src/python/pants/backend/project_info/peek.py

@@ -5,12 +5,14 @@
 
 import collections
 import json
+from abc import ABCMeta
 from dataclasses import dataclass, fields, is_dataclass
 from typing import Any, Iterable, Mapping, Protocol, runtime_checkable
 
 from pants.engine.addresses import Addresses
 from pants.engine.collection import Collection
 from pants.engine.console import Console
+from pants.engine.environment import EnvironmentName
 from pants.engine.fs import Snapshot
 from pants.engine.goal import Goal, GoalSubsystem, Outputting
 from pants.engine.internals.build_files import _get_target_family_and_adaptor_for_dep_rules
@@ -23,14 +25,18 @@
     DependenciesRuleApplication,
     DependenciesRuleApplicationRequest,
     Field,
+    FieldSet,
     HydratedSources,
     HydrateSourcesRequest,
+    ImmutableValue,
     SourcesField,
     Target,
     Targets,
     UnexpandedTargets,
 )
+from pants.engine.unions import UnionMembership, union
 from pants.option.option_types import BoolOption
+from pants.util.frozendict import FrozenDict
 from pants.util.strutil import softwrap
 
 
@@ -63,6 +69,10 @@ class PeekSubsystem(Outputting, GoalSubsystem):
         ),
     )
 
+    include_additional_info = BoolOption(
+        default=False, help="Whether to include additional information generated by plugins."
+    )
+
 
 class Peek(Goal):
     subsystem_cls = PeekSubsystem
@@ -75,6 +85,27 @@ def _normalize_value(val: Any) -> Any:
     return val
 
 
+@union(in_scope_types=[EnvironmentName])
+@dataclass(frozen=True)
+class HasAdditionalTargetDataFieldSet(FieldSet, metaclass=ABCMeta):
+    """Union type to attach data to a target that will appear under the "additional_info" field in
+    the output of `pants peek` if the `--peek-include-additional-info` option is enabled."""
+
+
+@dataclass(frozen=True)
+class AdditionalTargetData:
+    label: str
+    data: ImmutableValue
+
+    def __init__(self, label: str, data: Any) -> None:
+        object.__setattr__(self, "label", label)
+        if isinstance(data, collections.abc.Mapping):
+            data = FrozenDict.deep_freeze(data)
+        elif isinstance(data, (list, set)):
+            data = tuple(data)
+        object.__setattr__(self, "data", data)
+
+
 @dataclass(frozen=True)
 class TargetData:
     target: Target
@@ -85,6 +116,7 @@ class TargetData:
     dependencies_rules: tuple[str, ...] | None = None
     dependents_rules: tuple[str, ...] | None = None
     applicable_dep_rules: tuple[DependencyRuleApplication, ...] | None = None
+    additional_info: tuple[AdditionalTargetData, ...] | None = None
 
     def to_dict(self, exclude_defaults: bool = False, include_dep_rules: bool = False) -> dict:
         nothing = object()
@@ -106,6 +138,14 @@ def to_dict(self, exclude_defaults: bool = False, include_dep_rules: bool = Fals
             fields["_dependents_rules"] = self.dependents_rules
             fields["_applicable_dep_rules"] = self.applicable_dep_rules
 
+        if self.additional_info is not None:
+            fields["additional_info"] = {
+                additional_target_data.label: additional_target_data.data
+                for additional_target_data in sorted(
+                    self.additional_info, key=lambda atd: atd.label
+                )
+            }
+
         return {
             "address": self.target.address.spec,
             "target_type": self.target.alias,
@@ -165,6 +205,7 @@ async def get_target_data(
     # NB: We must preserve target generators, not replace with their generated targets.
     targets: UnexpandedTargets,
     subsys: PeekSubsystem,
+    union_membership: UnionMembership,
 ) -> TargetDatas:
     sorted_targets = sorted(targets, key=lambda tgt: tgt.address)
 
@@ -188,6 +229,26 @@ async def get_target_data(
         Get(HydratedSources, HydrateSourcesRequest(tgt[SourcesField]))
         for tgt in targets_with_sources
     )
+    if subsys.include_additional_info:
+        additional_info_field_sets = [
+            field_set_type.create(tgt)
+            for tgt in sorted_targets
+            for field_set_type in union_membership[HasAdditionalTargetDataFieldSet]
+            if field_set_type.is_applicable(tgt)
+        ]
+        additional_infos = await MultiGet(
+            Get(AdditionalTargetData, HasAdditionalTargetDataFieldSet, field_set)
+            for field_set in additional_info_field_sets
+        )
+        group_additional_infos_by_address_builder = collections.defaultdict(list)
+        for field_set, additional_info in zip(additional_info_field_sets, additional_infos):
+            group_additional_infos_by_address_builder[field_set.address].append(additional_info)
+        group_additional_infos_by_address = {
+            address: tuple(additional_info)
+            for address, additional_info in group_additional_infos_by_address_builder.items()
+        }
+    else:
+        group_additional_infos_by_address = {}
 
     expanded_dependencies = [
         tuple(dep.address.spec for dep in deps)
@@ -243,6 +304,9 @@ async def get_target_data(
             dependencies_rules=dependencies_rules_map.get(tgt.address),
             dependents_rules=dependents_rules_map.get(tgt.address),
             applicable_dep_rules=applicable_dep_rules_map.get(tgt.address),
+            additional_info=group_additional_infos_by_address.get(
+                tgt.address, () if subsys.include_additional_info else None
+            ),
         )
         for tgt, expanded_deps in zip(sorted_targets, expanded_dependencies)
     )

src/python/pants/backend/project_info/peek_test.py

@@ -5,26 +5,72 @@
 
 import dataclasses
 from textwrap import dedent
-from typing import Sequence
+from typing import Sequence, cast
 
 import pytest
 
 from pants.backend.project_info import peek
-from pants.backend.project_info.peek import Peek, TargetData, TargetDatas
+from pants.backend.project_info.peek import (
+    AdditionalTargetData,
+    HasAdditionalTargetDataFieldSet,
+    Peek,
+    TargetData,
+    TargetDatas,
+)
 from pants.backend.visibility.rules import rules as visibility_rules
 from pants.base.specs import RawSpecs, RecursiveGlobSpec
-from pants.core.target_types import ArchiveTarget, FilesGeneratorTarget, FileTarget, GenericTarget
+from pants.core.target_types import (
+    ArchiveTarget,
+    FilesGeneratorTarget,
+    FileSourceField,
+    FileTarget,
+    GenericTarget,
+)
 from pants.engine.addresses import Address
 from pants.engine.fs import Snapshot
 from pants.engine.internals.dep_rules import DependencyRuleAction, DependencyRuleApplication
-from pants.engine.rules import QueryRule
+from pants.engine.rules import QueryRule, rule
+from pants.engine.target import DescriptionField
+from pants.engine.unions import UnionRule
 from pants.testutil.rule_runner import RuleRunner
 
 
 def _snapshot(fingerprint: str, files: tuple[str, ...]) -> Snapshot:
     return Snapshot.create_for_testing(files, ())
 
 
+@dataclasses.dataclass(frozen=True)
+class FirstFakeAdditionalTargetDataFieldSet(HasAdditionalTargetDataFieldSet):
+    required_fields = (FileSourceField,)
+
+    source: FileSourceField
+
+
+@dataclasses.dataclass(frozen=True)
+class SecondFakeAdditionalTargetDataFieldSet(HasAdditionalTargetDataFieldSet):
+    required_fields = (FileSourceField, DescriptionField)
+
+    description: DescriptionField
+
+
+@rule
+async def first_fake_additional_target_data(
+    field_set: FirstFakeAdditionalTargetDataFieldSet,
+) -> AdditionalTargetData:
+    filename, extension = cast(str, field_set.source.value).split(".", 1)
+    return AdditionalTargetData("source_parts", {"filename": filename, "extension": extension})
+
+
+@rule
+async def second_fake_additional_target_data(
+    field_set: SecondFakeAdditionalTargetDataFieldSet,
+) -> AdditionalTargetData:
+    return AdditionalTargetData(
+        "reversed_description",
+        field_set.description.value[::-1] if field_set.description.value else None,
+    )
+
+
 @pytest.mark.parametrize(
     "expanded_target_infos, exclude_defaults, include_dep_rules, expected_output",
     [
@@ -249,6 +295,57 @@ def _snapshot(fingerprint: str, files: tuple[str, ...]) -> Snapshot:
             ),
             id="include-dep-rules",
         ),
+        pytest.param(
+            [
+                TargetData(
+                    FilesGeneratorTarget(
+                        {"sources": ["foo.txt"]}, Address("example", target_name="files_target")
+                    ),
+                    _snapshot(
+                        "1",
+                        ("foo.txt",),
+                    ),
+                    tuple(),
+                    additional_info=(
+                        AdditionalTargetData("test_data1", {"hello": "world"}),
+                        AdditionalTargetData("test_data2", ["one", "two"]),
+                    ),
+                )
+            ],
+            False,
+            False,
+            dedent(
+                """\
+                [
+                  {
+                    "address": "example:files_target",
+                    "target_type": "files",
+                    "additional_info": {
+                      "test_data1": {
+                        "hello": "world"
+                      },
+                      "test_data2": [
+                        "one",
+                        "two"
+                      ]
+                    },
+                    "dependencies": [],
+                    "description": null,
+                    "overrides": null,
+                    "sources": [
+                      "foo.txt"
+                    ],
+                    "sources_fingerprint": "b5e73bb1d7a3f8c2e7f8c43f38ab4d198e3512f082c670706df89f5abe319edf",
+                    "sources_raw": [
+                      "foo.txt"
+                    ],
+                    "tags": null
+                  }
+                ]
+                """
+            ),
+            id="include-additional-info",
+        ),
     ],
 )
 def test_render_targets_as_json(
@@ -264,9 +361,15 @@ def rule_runner() -> RuleRunner:
         rules=[
             *peek.rules(),
             *visibility_rules(),
+            UnionRule(HasAdditionalTargetDataFieldSet, FirstFakeAdditionalTargetDataFieldSet),
+            UnionRule(HasAdditionalTargetDataFieldSet, SecondFakeAdditionalTargetDataFieldSet),
+            first_fake_additional_target_data,
+            second_fake_additional_target_data,
             QueryRule(TargetDatas, [RawSpecs]),
+            QueryRule(AdditionalTargetData, [FirstFakeAdditionalTargetDataFieldSet]),
+            QueryRule(AdditionalTargetData, [SecondFakeAdditionalTargetDataFieldSet]),
         ],
-        target_types=[FilesGeneratorTarget, GenericTarget],
+        target_types=[FilesGeneratorTarget, GenericTarget, FileTarget],
     )
 
 
@@ -392,3 +495,37 @@ def test_get_target_data_with_dep_rules(rule_runner: RuleRunner) -> None:
             applicable_dep_rules=(),
         ),
     ]
+
+
+def test_get_target_data_with_additional_info(rule_runner: RuleRunner) -> None:
+    rule_runner.set_options(["--peek-include-additional-info"])
+    rule_runner.write_files(
+        {
+            "foo/BUILD": dedent(
+                """\
+            file(source="a.txt", description="reverse me!")
+            """
+            ),
+            "foo/a.txt": "",
+        }
+    )
+    tds = rule_runner.request(
+        TargetDatas,
+        [RawSpecs(recursive_globs=(RecursiveGlobSpec("foo"),), description_of_origin="tests")],
+    )
+
+    assert _normalize_fingerprints(tds) == [
+        TargetData(
+            FileTarget(
+                {"source": "a.txt", "description": "reverse me!"},
+                Address("foo", target_name="foo"),
+                name_explicitly_set=False,
+            ),
+            _snapshot("", ("foo/a.txt",)),
+            (),
+            additional_info=(
+                AdditionalTargetData("source_parts", {"filename": "a", "extension": "txt"}),
+                AdditionalTargetData("reversed_description", "reverse me!"[::-1]),
+            ),
+        ),
+    ]