Extend dependents goal with output format to support JSON

Author: AlexTereshenkov

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

@@ -1,8 +1,9 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
-
+import json
 from collections import defaultdict
 from dataclasses import dataclass
+from enum import Enum
 from typing import Iterable, Set
 
 from pants.engine.addresses import Address, Addresses
@@ -16,7 +17,7 @@
     Dependencies,
     DependenciesRequest,
 )
-from pants.option.option_types import BoolOption
+from pants.option.option_types import BoolOption, EnumOption
 from pants.util.frozendict import FrozenDict
 from pants.util.logging import LogLevel
 from pants.util.ordered_set import FrozenOrderedSet
@@ -27,6 +28,17 @@ class AddressToDependents:
     mapping: FrozenDict[Address, FrozenOrderedSet[Address]]
 
 
+class DependentsOutputFormat(Enum):
+    """Output format for listing dependents.
+
+    text: List all dependents as a single list of targets in plain text.
+    json: List all dependents as a mapping `{target: [dependents]}`.
+    """
+
+    text = "text"
+    json = "json"
+
+
 @rule(desc="Map all targets to their dependents", level=LogLevel.DEBUG)
 async def map_addresses_to_dependents(all_targets: AllUnexpandedTargets) -> AddressToDependents:
     dependencies_per_target = await MultiGet(
@@ -105,7 +117,12 @@ class DependentsSubsystem(LineOriented, GoalSubsystem):
     )
     closed = BoolOption(
         default=False,
-        help="Include the input targets in the output, along with the dependents.",
+        help="Include the input targets in the output, along with the dependents. This option "
+        "only applies when using the `text` format.",
+    )
+    format = EnumOption(
+        default=DependentsOutputFormat.text,
+        help="Output format for listing dependents.",
     )
 
 
@@ -114,21 +131,66 @@ class DependentsGoal(Goal):
     environment_behavior = Goal.EnvironmentBehavior.LOCAL_ONLY
 
 
-@goal_rule
-async def dependents_goal(
-    specified_addresses: Addresses, dependents_subsystem: DependentsSubsystem, console: Console
-) -> DependentsGoal:
+async def list_dependents_as_plain_text(
+    addresses: Addresses, dependents_subsystem: DependentsSubsystem, console: Console
+) -> None:
+    """Get dependents for given addresses and list them in the console as a single list."""
     dependents = await Get(
         Dependents,
         DependentsRequest(
-            specified_addresses,
+            addresses,
             transitive=dependents_subsystem.transitive,
             include_roots=dependents_subsystem.closed,
         ),
     )
     with dependents_subsystem.line_oriented(console) as print_stdout:
         for address in dependents:
             print_stdout(address.spec)
+
+
+async def list_dependents_as_json(
+    addresses: Addresses, dependents_subsystem: DependentsSubsystem, console: Console
+) -> None:
+    """Get dependents for given addresses and list them in the console in JSON.
+
+    Note that `--closed` option is ignored as it doesn't make sense to duplicate source address in
+    the list of its dependents.
+    """
+    dependents_group = await MultiGet(
+        Get(
+            Dependents,
+            DependentsRequest(
+                (address,),
+                transitive=dependents_subsystem.transitive,
+                include_roots=False,
+            ),
+        )
+        for address in addresses
+    )
+    iterated_addresses = []
+    for dependents in dependents_group:
+        iterated_addresses.append(sorted([str(address) for address in dependents]))
+    mapping = dict(zip([str(address) for address in addresses], iterated_addresses))
+    output = json.dumps(mapping, indent=4)
+    console.print_stdout(output)
+
+
+@goal_rule
+async def dependents_goal(
+    specified_addresses: Addresses, dependents_subsystem: DependentsSubsystem, console: Console
+) -> DependentsGoal:
+    if DependentsOutputFormat.text == dependents_subsystem.format:
+        await list_dependents_as_plain_text(
+            addresses=specified_addresses,
+            dependents_subsystem=dependents_subsystem,
+            console=console,
+        )
+    elif DependentsOutputFormat.json == dependents_subsystem.format:
+        await list_dependents_as_json(
+            addresses=specified_addresses,
+            dependents_subsystem=dependents_subsystem,
+            console=console,
+        )
     return DependentsGoal(exit_code=0)
 
 

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

@@ -1,11 +1,12 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
-
-from typing import List
+import json
+from functools import partial
+from typing import Dict, List, Union
 
 import pytest
 
-from pants.backend.project_info.dependents import DependentsGoal
+from pants.backend.project_info.dependents import DependentsGoal, DependentsOutputFormat
 from pants.backend.project_info.dependents import rules as dependent_rules
 from pants.engine.target import Dependencies, SpecialCasedDependencies, Target
 from pants.testutil.rule_runner import RuleRunner
@@ -41,17 +42,22 @@ def assert_dependents(
     rule_runner: RuleRunner,
     *,
     targets: List[str],
-    expected: List[str],
+    expected: Union[List[str], Dict[str, List[str]]],
     transitive: bool = False,
     closed: bool = False,
+    output_format: DependentsOutputFormat = DependentsOutputFormat.text,
 ) -> None:
     args = []
     if transitive:
         args.append("--transitive")
     if closed:
         args.append("--closed")
+    args.append(f"--format={output_format.value}")
     result = rule_runner.run_goal_rule(DependentsGoal, args=[*args, *targets])
-    assert result.stdout.splitlines() == expected
+    if output_format == DependentsOutputFormat.text:
+        assert result.stdout.splitlines() == expected
+    elif output_format == DependentsOutputFormat.json:
+        assert json.loads(result.stdout) == expected
 
 
 def test_no_targets(rule_runner: RuleRunner) -> None:
@@ -102,3 +108,82 @@ def test_special_cased_dependencies(rule_runner: RuleRunner) -> None:
         transitive=True,
         expected=["intermediate:intermediate", "leaf:leaf", "special:special"],
     )
+
+
+def test_dependents_as_json_direct_deps(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files({"special/BUILD": "tgt(special_deps=['intermediate'])"})
+    assert_deps = partial(
+        assert_dependents,
+        rule_runner,
+        output_format=DependentsOutputFormat.json,
+    )
+    # input: single target
+    assert_deps(
+        targets=["base"],
+        transitive=False,
+        expected={
+            "base:base": ["intermediate:intermediate"],
+        },
+    )
+
+    # input: multiple targets
+    assert_deps(
+        targets=["base", "intermediate"],
+        transitive=False,
+        expected={
+            "base:base": ["intermediate:intermediate"],
+            "intermediate:intermediate": ["leaf:leaf", "special:special"],
+        },
+    )
+
+    # input: all targets
+    assert_deps(
+        targets=["::"],
+        transitive=False,
+        expected={
+            "base:base": ["intermediate:intermediate"],
+            "intermediate:intermediate": ["leaf:leaf", "special:special"],
+            "leaf:leaf": [],
+            "special:special": [],
+        },
+    )
+
+
+def test_dependents_as_json_transitive_deps(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files({"special/BUILD": "tgt(special_deps=['intermediate'])"})
+    assert_deps = partial(
+        assert_dependents,
+        rule_runner,
+        output_format=DependentsOutputFormat.json,
+    )
+
+    # input: single target
+    assert_deps(
+        targets=["base"],
+        transitive=True,
+        expected={
+            "base:base": ["intermediate:intermediate", "leaf:leaf", "special:special"],
+        },
+    )
+
+    # input: multiple targets
+    assert_deps(
+        targets=["base", "intermediate"],
+        transitive=True,
+        expected={
+            "base:base": ["intermediate:intermediate", "leaf:leaf", "special:special"],
+            "intermediate:intermediate": ["leaf:leaf", "special:special"],
+        },
+    )
+
+    # input: all targets
+    assert_deps(
+        targets=["::"],
+        transitive=False,
+        expected={
+            "base:base": ["intermediate:intermediate"],
+            "intermediate:intermediate": ["leaf:leaf", "special:special"],
+            "leaf:leaf": [],
+            "special:special": [],
+        },
+    )