More flexible binary and run rules.

src/python/pants/rules/core/binary.py

@@ -6,18 +6,17 @@
 import os
 from abc import ABC
 from dataclasses import dataclass
-from typing import ClassVar, Dict, Iterable, Tuple, Type
+from typing import ClassVar, Dict, Iterable, List, Mapping, Sequence, Tuple, Type
 
 from pants.base.build_root import BuildRoot
 from pants.build_graph.address import Address
-from pants.engine.addressable import Addresses
 from pants.engine.console import Console
 from pants.engine.fs import Digest, DirectoriesToMerge, DirectoryToMaterialize, Workspace
 from pants.engine.goal import Goal, GoalSubsystem, LineOriented
 from pants.engine.objects import union
-from pants.engine.rules import UnionMembership, goal_rule, rule
+from pants.engine.rules import UnionMembership, goal_rule
 from pants.engine.selectors import Get, MultiGet
-from pants.engine.target import Field, RegisteredTargetTypes, Target, WrappedTarget
+from pants.engine.target import Field, RegisteredTargetTypes, Target, TargetsWithOrigins
 from pants.rules.core.distdir import DistDir
 
 
@@ -84,74 +83,91 @@ class Binary(Goal):
     subsystem_cls = BinaryOptions
 
 
-@goal_rule
-async def create_binary(
-    addresses: Addresses,
-    console: Console,
-    workspace: Workspace,
-    options: BinaryOptions,
-    distdir: DistDir,
-    buildroot: BuildRoot,
-) -> Binary:
-    with options.line_oriented(console) as print_stdout:
-        binaries = await MultiGet(Get[CreatedBinary](Address, address) for address in addresses)
-        merged_digest = await Get[Digest](
-            DirectoriesToMerge(tuple(binary.digest for binary in binaries))
-        )
-        result = workspace.materialize_directory(
-            DirectoryToMaterialize(merged_digest, path_prefix=str(distdir.relpath))
-        )
-        for path in result.output_paths:
-            print_stdout(f"Wrote {os.path.relpath(path, buildroot.path)}")
-    return Binary(exit_code=0)
-
-
-@rule
-async def coordinator_of_binaries(
-    wrapped_target: WrappedTarget,
+def gather_valid_binary_configuration_types(
+    goal_subsytem: GoalSubsystem,
+    targets_with_origins: TargetsWithOrigins,
     union_membership: UnionMembership,
     registered_target_types: RegisteredTargetTypes,
-) -> CreatedBinary:
-    target = wrapped_target.target
+) -> Mapping[Target, Sequence[Type[BinaryConfiguration]]]:
     config_types: Iterable[Type[BinaryConfiguration]] = union_membership.union_rules[
         BinaryConfiguration
     ]
-    valid_config_types = [
-        config_type for config_type in config_types if config_type.is_valid(target)
-    ]
-    if not valid_config_types:
+    valid_config_types_by_target: Dict[Target, List[Type[BinaryConfiguration]]] = {}
+    valid_target_aliases = set()
+    for target_with_origin in targets_with_origins:
+        for config_type in config_types:
+            if config_type.is_valid(target_with_origin.target):
+                valid_config_types_by_target.setdefault(target_with_origin.target, []).append(
+                    config_type
+                )
+                valid_target_aliases.add(target_with_origin.target.alias)
+
+    if not valid_config_types_by_target:
         all_valid_target_types = itertools.chain.from_iterable(
             config_type.valid_target_types(
                 registered_target_types.types, union_membership=union_membership
             )
             for config_type in config_types
         )
-        formatted_target_types = sorted(target_type.alias for target_type in all_valid_target_types)
-        # TODO: this is a leaky abstraction that the error message knows this rule is being used
-        #  by `run` and `binary`. How should this be handled? A custom goal author could depend on
-        #  this error message too and the error would now be lying.
-        raise ValueError(
-            f"The `run` and `binary` goals only work with the following target types: "
-            f"{formatted_target_types}\n\nYou used {target.address} with target type "
-            f"{repr(target.alias)}."
+        all_valid_target_aliases = sorted(
+            target_type.alias for target_type in all_valid_target_types
+        )
+        invalid_target_aliases = sorted(
+            {
+                target_with_origin.target.alias
+                for target_with_origin in targets_with_origins
+                if target_with_origin.target.alias not in valid_target_aliases
+            }
+        )
+        specs = sorted(
+            {
+                target_with_origin.origin.to_spec_string()
+                for target_with_origin in targets_with_origins
+            }
         )
-    # TODO: we must use this check when running `./v2 run` because we should only run one target
-    #  with one implementation. But, we don't necessarily need to enforce this with `./v2 binary`.
-    #  See https://github.com/pantsbuild/pants/pull/9345#discussion_r395221542 for some possible
-    #  semantics for `./v2 binary`.
-    if len(valid_config_types) > 1:
-        possible_config_types = sorted(config_type.__name__ for config_type in valid_config_types)
-        # TODO: improve this error message. (It's never actually triggered yet because we only have
-        #  Python implemented with V2.) A better error message would explain to users how they can
-        #  resolve the issue.
+        bulleted_list_sep = "\n  * "
         raise ValueError(
-            f"Multiple of the registered binary implementations work for {target.address} "
-            f"(target type {repr(target.alias)}). It is ambiguous which implementation to use. "
-            f"Possible implementations: {possible_config_types}."
+            f"The `{goal_subsytem.name}` goal only works with the following target types:"
+            f"{bulleted_list_sep}{bulleted_list_sep.join(all_valid_target_aliases)}\n\n"
+            f"You specified `{' '.join(specs)}` which only included the following target types:"
+            f"{bulleted_list_sep}{bulleted_list_sep.join(invalid_target_aliases)}"
         )
-    config_type = valid_config_types[0]
-    return await Get[CreatedBinary](BinaryConfiguration, config_type.create(target))
+    return valid_config_types_by_target
+
+
+@goal_rule
+async def create_binary(
+    targets_with_origins: TargetsWithOrigins,
+    console: Console,
+    workspace: Workspace,
+    options: BinaryOptions,
+    distdir: DistDir,
+    buildroot: BuildRoot,
+    union_membership: UnionMembership,
+    registered_target_types: RegisteredTargetTypes,
+) -> Binary:
+    valid_config_types_by_target = gather_valid_binary_configuration_types(
+        goal_subsytem=options,
+        targets_with_origins=targets_with_origins,
+        union_membership=union_membership,
+        registered_target_types=registered_target_types,
+    )
+    binaries = await MultiGet(
+        Get[CreatedBinary](BinaryConfiguration, valid_config_type.create(target))
+        for target, valid_config_types in valid_config_types_by_target.items()
+        for valid_config_type in valid_config_types
+    )
+    merged_digest = await Get[Digest](
+        DirectoriesToMerge(tuple(binary.digest for binary in binaries))
+    )
+    result = workspace.materialize_directory(
+        DirectoryToMaterialize(merged_digest, path_prefix=str(distdir.relpath))
+    )
+    with options.line_oriented(console) as print_stdout:
+        for path in result.output_paths:
+            print_stdout(f"Wrote {os.path.relpath(path, buildroot.path)}")
+    return Binary(exit_code=0)
 
 
 def rules():
-    return [create_binary, coordinator_of_binaries]
+    return [create_binary]

src/python/pants/rules/core/run.py

@@ -4,17 +4,20 @@
 from pathlib import PurePath
 
 from pants.base.build_root import BuildRoot
-from pants.build_graph.address import Address
-from pants.engine.addressable import Addresses
 from pants.engine.console import Console
 from pants.engine.fs import DirectoryToMaterialize, Workspace
 from pants.engine.goal import Goal, GoalSubsystem
 from pants.engine.interactive_runner import InteractiveProcessRequest, InteractiveRunner
-from pants.engine.rules import goal_rule
+from pants.engine.rules import UnionMembership, goal_rule
 from pants.engine.selectors import Get
+from pants.engine.target import RegisteredTargetTypes, TargetsWithOrigins
 from pants.option.custom_types import shell_str
 from pants.option.global_options import GlobalOptions
-from pants.rules.core.binary import BinaryConfiguration, CreatedBinary
+from pants.rules.core.binary import (
+    BinaryConfiguration,
+    CreatedBinary,
+    gather_valid_binary_configuration_types,
+)
 from pants.util.contextutil import temporary_dir
 
 
@@ -49,12 +52,47 @@ async def run(
     workspace: Workspace,
     runner: InteractiveRunner,
     build_root: BuildRoot,
-    addresses: Addresses,
+    targets_with_origins: TargetsWithOrigins,
     options: RunOptions,
     global_options: GlobalOptions,
+    union_membership: UnionMembership,
+    registered_target_types: RegisteredTargetTypes,
 ) -> Run:
-    address = addresses.expect_single()
-    binary = await Get[CreatedBinary](Address, address)
+    valid_config_types_by_target = gather_valid_binary_configuration_types(
+        goal_subsytem=options,
+        targets_with_origins=targets_with_origins,
+        union_membership=union_membership,
+        registered_target_types=registered_target_types,
+    )
+
+    bulleted_list_sep = "\n  * "
+
+    if len(valid_config_types_by_target) > 1:
+        binary_target_addresses = sorted(
+            binary_target.address.spec for binary_target in valid_config_types_by_target
+        )
+        raise ValueError(
+            f"The `run` goal only works on one binary target but was given multiple targets that "
+            f"can produce a binary:"
+            f"{bulleted_list_sep}{bulleted_list_sep.join(binary_target_addresses)}\n\n"
+            f"Please select one of these targets to run."
+        )
+
+    target, valid_config_types = list(valid_config_types_by_target.items())[0]
+    if len(valid_config_types) > 1:
+        possible_config_types = sorted(config_type.__name__ for config_type in valid_config_types)
+        # TODO: improve this error message. (It's never actually triggered yet because we only have
+        #  Python implemented with V2.) A better error message would explain to users how they can
+        #  resolve the issue.
+        raise ValueError(
+            f"Multiple of the registered binary implementations work for {target.address} "
+            f"(target type {repr(target.alias)}).\n\n"
+            f"It is ambiguous which implementation to use. Possible implementations:"
+            f"{bulleted_list_sep}{bulleted_list_sep.join(possible_config_types)}."
+        )
+    config_type = valid_config_types[0]
+
+    binary = await Get[CreatedBinary](BinaryConfiguration, config_type.create(target))
 
     workdir = global_options.options.pants_workdir
 
@@ -64,7 +102,7 @@ async def run(
             DirectoryToMaterialize(binary.digest, path_prefix=path_relative_to_build_root)
         )
 
-        console.write_stdout(f"Running target: {address}\n")
+        console.write_stdout(f"Running target: {target.address}\n")
         full_path = PurePath(tmpdir, binary.binary_name).as_posix()
         run_request = InteractiveProcessRequest(
             argv=(full_path, *options.values.args), run_in_workspace=True,
@@ -74,12 +112,14 @@ async def run(
             result = runner.run_local_interactive_process(run_request)
             exit_code = result.process_exit_code
             if result.process_exit_code == 0:
-                console.write_stdout(f"{address} ran successfully.\n")
+                console.write_stdout(f"{target.address} ran successfully.\n")
             else:
-                console.write_stderr(f"{address} failed with code {result.process_exit_code}!\n")
+                console.write_stderr(
+                    f"{target.address} failed with code {result.process_exit_code}!\n"
+                )
 
         except Exception as e:
-            console.write_stderr(f"Exception when attempting to run {address}: {e!r}\n")
+            console.write_stderr(f"Exception when attempting to run {target.address}: {e!r}\n")
             exit_code = -1
 
     return Run(exit_code)

src/python/pants/rules/core/run_test.py

@@ -4,12 +4,14 @@
 from typing import cast
 
 from pants.base.build_root import BuildRoot
+from pants.base.specs import SingleAddress
 from pants.build_graph.address import Address
-from pants.engine.addressable import Addresses
 from pants.engine.fs import Digest, FileContent, InputFilesContent, Workspace
 from pants.engine.interactive_runner import InteractiveProcessRequest, InteractiveRunner
+from pants.engine.rules import UnionMembership
+from pants.engine.target import RegisteredTargetTypes, Target, TargetsWithOrigins, TargetWithOrigin
 from pants.option.global_options import GlobalOptions
-from pants.rules.core.binary import CreatedBinary
+from pants.rules.core.binary import BinaryConfiguration, CreatedBinary
 from pants.rules.core.run import Run, RunOptions, run
 from pants.testutil.engine.util import (
     MockConsole,
@@ -19,6 +21,7 @@
     run_rule,
 )
 from pants.testutil.test_base import TestBase
+from pants.util.ordered_set import OrderedSet
 
 
 class RunTest(TestBase):
@@ -34,22 +37,42 @@ def single_target_run(
     ) -> Run:
         workspace = Workspace(self.scheduler)
         interactive_runner = InteractiveRunner(self.scheduler)
-        BuildRoot().path = self.build_root
+
+        class TestBinaryConfiguration(BinaryConfiguration):
+            required_fields = ()
+
+        class TestBinaryTarget(Target):
+            alias = "binary"
+            core_fields = ()
+
+        address = Address.parse(address_spec)
+        origin = SingleAddress(address.spec_path, address.target_name)
         res = run_rule(
             run,
             rule_args=[
                 console,
                 workspace,
                 interactive_runner,
                 BuildRoot(),
-                Addresses([Address.parse(address_spec)]),
+                TargetsWithOrigins(
+                    [
+                        TargetWithOrigin(
+                            target=TestBinaryTarget(unhydrated_values={}, address=address),
+                            origin=origin,
+                        )
+                    ]
+                ),
                 create_goal_subsystem(RunOptions, args=[]),
                 create_subsystem(GlobalOptions, pants_workdir=self.pants_workdir),
+                UnionMembership(
+                    union_rules={BinaryConfiguration: OrderedSet([TestBinaryConfiguration])}
+                ),
+                RegisteredTargetTypes.create([TestBinaryTarget]),
             ],
             mock_gets=[
                 MockGet(
                     product_type=CreatedBinary,
-                    subject_type=Address,
+                    subject_type=TestBinaryConfiguration,
                     mock=lambda _: self.create_mock_binary(program_text),
                 ),
             ],