Pre-filter TestTarget @union members

src/python/pants/build_graph/build_configuration.py

@@ -170,7 +170,7 @@ def rules(self):
     return list(self._rules)
 
   def union_rules(self):
-    """Returns a mapping of registered union base types -> [a list of union member types].
+    """Returns a mapping of registered union base types -> [OrderedSet of union member types].
 
     :rtype: OrderedDict
     """

src/python/pants/engine/rules.py

@@ -6,12 +6,12 @@
 import itertools
 import logging
 import sys
-import typing
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from collections.abc import Iterable
+from dataclasses import dataclass
 from textwrap import dedent
-from typing import Any, Callable, Type, cast
+from typing import Any, Callable, Dict, Iterable, Type, cast
 
 import asttokens
 from twitter.common.collections import OrderedSet
@@ -211,7 +211,7 @@ def _get_starting_indent(source):
 
 
 def _make_rule(
-  return_type: type, parameter_types: typing.Iterable[type], cacheable: bool = True
+  return_type: type, parameter_types: Iterable[type], cacheable: bool = True
 ) -> Callable[[Callable], Callable]:
   """A @decorator that declares that a particular static function may be used as a TaskRule.
 
@@ -428,6 +428,17 @@ def __new__(cls, union_base, union_member):
     return super().__new__(cls, union_base, union_member)
 
 
+@dataclass(frozen=True)
+class UnionMembership:
+  union_rules: Dict[type, Iterable[type]]
+
+  def is_member(self, union_type, putative_member):
+    members = self.union_rules.get(union_type)
+    if members is None:
+      raise TypeError(f'Not a registered union type: {union_type}')
+    return type(putative_member) in members
+
+
 class Rule(ABC):
   """Rules declare how to produce products for the product graph.
 

src/python/pants/init/engine_initializer.py

@@ -42,7 +42,7 @@
 from pants.engine.mapper import AddressMapper
 from pants.engine.parser import SymbolTable
 from pants.engine.platform import create_platform_rules
-from pants.engine.rules import RootRule, rule
+from pants.engine.rules import RootRule, UnionMembership, rule
 from pants.engine.scheduler import Scheduler
 from pants.engine.selectors import Params
 from pants.init.options_initializer import BuildConfigInitializer, OptionsInitializer
@@ -357,6 +357,10 @@ def build_configuration_singleton() -> BuildConfiguration:
     def symbol_table_singleton() -> SymbolTable:
       return symbol_table
 
+    @rule
+    def union_membership_singleton() -> UnionMembership:
+      return UnionMembership(build_configuration.union_rules())
+
     # Create a Scheduler containing graph and filesystem rules, with no installed goals. The
     # LegacyBuildGraph will explicitly request the products it needs.
     rules = (
@@ -365,6 +369,7 @@ def symbol_table_singleton() -> SymbolTable:
         glob_match_error_behavior_singleton,
         build_configuration_singleton,
         symbol_table_singleton,
+        union_membership_singleton,
       ] +
       create_legacy_graph_tasks() +
       create_fs_rules() +

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

@@ -4,12 +4,10 @@
 import logging
 
 from pants.base.exiter import PANTS_FAILED_EXIT_CODE, PANTS_SUCCEEDED_EXIT_CODE
-from pants.build_graph.address import Address
-from pants.engine.addressable import BuildFileAddresses
 from pants.engine.console import Console
 from pants.engine.goal import Goal
-from pants.engine.legacy.graph import HydratedTarget
-from pants.engine.rules import console_rule, rule
+from pants.engine.legacy.graph import HydratedTarget, HydratedTargets
+from pants.engine.rules import UnionMembership, console_rule, rule
 from pants.engine.selectors import Get
 from pants.rules.core.core_test_model import Status, TestResult, TestTarget
 
@@ -25,33 +23,39 @@ class Test(Goal):
 
 
 @console_rule
-def fast_test(console: Console, addresses: BuildFileAddresses) -> Test:
-  test_results = yield [Get(TestResult, Address, address.to_address()) for address in addresses]
+def fast_test(console: Console,
+              targets: HydratedTargets,
+              union_membership: UnionMembership) -> Test:
+  filtered_targets = [tgt for tgt in targets if union_membership.is_member(TestTarget, tgt.adaptor)]
+  test_results = yield [Get(TestResult, HydratedTarget, tgt) for tgt in filtered_targets]
   did_any_fail = False
-  for address, test_result in zip(addresses, test_results):
+  for tgt, test_result in zip(filtered_targets, test_results):
     if test_result.status == Status.FAILURE:
       did_any_fail = True
     if test_result.stdout:
       console.write_stdout(
         "{} stdout:\n{}\n".format(
-          address.reference(),
-          console.red(test_result.stdout) if test_result.status == Status.FAILURE else test_result.stdout
+          tgt.address.reference(),
+          (console.red(test_result.stdout) if test_result.status == Status.FAILURE
+           else test_result.stdout)
         )
       )
     if test_result.stderr:
       # NB: we write to stdout, rather than to stderr, to avoid potential issues interleaving the
       # two streams.
       console.write_stdout(
         "{} stderr:\n{}\n".format(
-          address.reference(),
-          console.red(test_result.stderr) if test_result.status == Status.FAILURE else test_result.stderr
+          tgt.address.reference(),
+          (console.red(test_result.stderr) if test_result.status == Status.FAILURE
+           else test_result.stderr)
         )
       )
 
   console.write_stdout("\n")
 
-  for address, test_result in zip(addresses, test_results):
-    console.print_stdout('{0:80}.....{1:>10}'.format(address.reference(), test_result.status.value))
+  for tgt, test_result in zip(filtered_targets, test_results):
+    console.print_stdout('{0:80}.....{1:>10}'.format(
+      tgt.address.reference(), test_result.status.value))
 
   if did_any_fail:
     console.print_stderr(console.red('Tests failed'))