improve ergonomics of checking for @union types with new @decorated_t…

…ype_checkable decorator (pantsbuild#8496)

### Problem

We manually check `getattr(some_type, '_is_union', False)` in multiple places in order to check if `some_type` was declared with the `@union` annotation. We would like to canonicalize this check and make it type-safe in some sense.

### Solution

- Create `@decorated_type_checkable` in `pants.util.meta`, which contains two methods:
  1. `define_instance_of(cls)`, which wraps a class `cls` and sets the attribute `_decorated_type_checkable_type` to the return value of `@decorated_type_checkable`.
  2. `is_instance(cls)`, which checks whether the sentinel attribute `_decorated_type_checkable_type` matches the type of the return value of `@decorated_type_checkable`.

src/python/pants/backend/awslambda/common/awslambda_common_rules.py

@@ -9,7 +9,8 @@
 from pants.engine.fs import Digest, DirectoriesToMerge, DirectoryToMaterialize, Workspace
 from pants.engine.goal import Goal, GoalSubsystem, LineOriented
 from pants.engine.legacy.graph import HydratedTarget
-from pants.engine.rules import console_rule, rule, union
+from pants.engine.objects import union
+from pants.engine.rules import console_rule, rule
 from pants.engine.selectors import Get, MultiGet
 from pants.rules.core.distdir import DistDir
 

src/python/pants/backend/python/lint/python_format_target.py

@@ -12,7 +12,8 @@
   PythonTestsAdaptor,
   TargetAdaptor,
 )
-from pants.engine.rules import UnionMembership, UnionRule, rule, union
+from pants.engine.objects import union
+from pants.engine.rules import UnionMembership, UnionRule, rule
 from pants.engine.selectors import Get
 from pants.rules.core.fmt import AggregatedFmtResults, FmtResult, FormatTarget
 

src/python/pants/backend/python/lint/python_lint_target.py

@@ -11,7 +11,8 @@
   PythonTestsAdaptor,
   TargetAdaptor,
 )
-from pants.engine.rules import UnionMembership, UnionRule, rule, union
+from pants.engine.objects import union
+from pants.engine.rules import UnionMembership, UnionRule, rule
 from pants.engine.selectors import Get, MultiGet
 from pants.rules.core.lint import LintResult, LintResults, LintTarget
 

src/python/pants/engine/legacy/structs.py

@@ -11,8 +11,8 @@
 from pants.build_graph.target import Target
 from pants.engine.addressable import addressable_list
 from pants.engine.fs import GlobExpansionConjunction, PathGlobs
-from pants.engine.objects import Locatable
-from pants.engine.rules import UnionRule, union
+from pants.engine.objects import Locatable, union
+from pants.engine.rules import UnionRule
 from pants.engine.struct import Struct, StructWithDeps
 from pants.source import wrapped_globs
 from pants.util.contextutil import exception_logging

src/python/pants/engine/native.py

@@ -37,6 +37,7 @@
   FallibleExecuteProcessResult,
   MultiPlatformExecuteProcessRequest,
 )
+from pants.engine.objects import union
 from pants.engine.selectors import Get
 from pants.util.contextutil import temporary_dir
 from pants.util.dirutil import read_file, safe_mkdir, safe_mkdtemp
@@ -315,7 +316,7 @@ def extern_is_union(self, context_handle, type_id):
     """Return whether or not a type is a member of a union"""
     c = self._ffi.from_handle(context_handle)
     input_type = c.from_id(type_id.tup_0)
-    return bool(getattr(input_type, '_is_union', None))
+    return union.is_instance(input_type)
 
   _do_raise_keyboardinterrupt_on_identify = bool(os.environ.get('_RAISE_KEYBOARDINTERRUPT_IN_CFFI_IDENTIFY', False))
 

src/python/pants/engine/objects.py

@@ -8,6 +8,8 @@
 from collections.abc import Iterable
 from typing import Generic, Iterator, TypeVar
 
+from pants.util.meta import decorated_type_checkable
+
 
 class SerializationError(Exception):
   """Indicates an error serializing an object."""
@@ -167,3 +169,41 @@ def __iter__(self) -> Iterator[_C]:
 
   def __bool__(self) -> bool:
     return bool(self.dependencies)
+
+
+@decorated_type_checkable
+def union(cls):
+  """A class decorator which other classes can specify that they can resolve to with `UnionRule`.
+
+  Annotating a class with @union allows other classes to use a UnionRule() instance to indicate that
+  they can be resolved to this base union class. This class will never be instantiated, and should
+  have no members -- it is used as a tag only, and will be replaced with whatever object is passed
+  in as the subject of a `await Get(...)`. See the following example:
+
+  @union
+  class UnionBase: pass
+
+  @rule
+  async def get_some_union_type(x: X) -> B:
+    result = await Get(ResultType, UnionBase, x.f())
+    # ...
+
+  If there exists a single path from (whatever type the expression `x.f()` returns) -> `ResultType`
+  in the rule graph, the engine will retrieve and execute that path to produce a `ResultType` from
+  `x.f()`. This requires also that whatever type `x.f()` returns was registered as a union member of
+  `UnionBase` with a `UnionRule`.
+
+  Unions allow @rule bodies to be written without knowledge of what types may eventually be provided
+  as input -- rather, they let the engine check that there is a valid path to the desired result.
+  """
+  # TODO: Check that the union base type is used as a tag and nothing else (e.g. no attributes)!
+  assert isinstance(cls, type)
+  def non_member_error_message(subject):
+    if hasattr(cls, 'non_member_error_message'):
+      return cls.non_member_error_message(subject)
+    desc = f' ("{cls.__doc__}")' if cls.__doc__ else ''
+    return f'Type {type(subject).__name__} is not a member of the {cls.__name__} @union{desc}'
+
+  return union.define_instance_of(
+    cls,
+    non_member_error_message=staticmethod(non_member_error_message))

src/python/pants/engine/rules.py

@@ -14,6 +14,7 @@
 from twitter.common.collections import OrderedSet
 
 from pants.engine.goal import Goal
+from pants.engine.objects import union
 from pants.engine.selectors import Get
 from pants.util.collections import assert_single_element
 from pants.util.memo import memoized
@@ -250,52 +251,14 @@ def console_rule(*args, **kwargs) -> Callable:
   return inner_rule(*args, **kwargs, cacheable=False)
 
 
-def union(cls):
-  """A class decorator which other classes can specify that they can resolve to with `UnionRule`.
-
-  Annotating a class with @union allows other classes to use a UnionRule() instance to indicate that
-  they can be resolved to this base union class. This class will never be instantiated, and should
-  have no members -- it is used as a tag only, and will be replaced with whatever object is passed
-  in as the subject of a `await Get(...)`. See the following example:
-
-  @union
-  class UnionBase: pass
-
-  @rule
-  def get_some_union_type(x: X) -> B:
-    result = await Get(ResultType, UnionBase, x.f())
-    # ...
-
-  If there exists a single path from (whatever type the expression `x.f()` returns) -> `ResultType`
-  in the rule graph, the engine will retrieve and execute that path to produce a `ResultType` from
-  `x.f()`. This requires also that whatever type `x.f()` returns was registered as a union member of
-  `UnionBase` with a `UnionRule`.
-
-  Unions allow @rule bodies to be written without knowledge of what types may eventually be provided
-  as input -- rather, they let the engine check that there is a valid path to the desired result.
-  """
-  # TODO: Check that the union base type is used as a tag and nothing else (e.g. no attributes)!
-  assert isinstance(cls, type)
-  def non_member_error_message(subject):
-    if hasattr(cls, 'non_member_error_message'):
-      return cls.non_member_error_message(subject)
-    desc = f' ("{cls.__doc__}")' if cls.__doc__ else ''
-    return f'Type {type(subject).__name__} is not a member of the {cls.__name__} @union{desc}'
-
-  return type(cls.__name__, (cls,), {
-    '_is_union': True,
-    'non_member_error_message': non_member_error_message,
-  })
-
-
 @dataclass(frozen=True)
 class UnionRule:
   """Specify that an instance of `union_member` can be substituted wherever `union_base` is used."""
   union_base: Type
   union_member: Type
 
   def __post_init__(self) -> None:
-    if not getattr(self.union_base, '_is_union', False):
+    if not union.is_instance(self.union_base):
       raise ValueError(
         f'union_base must be a type annotated with @union: was {self.union_base} '
         f'(type {type(self.union_base).__name__})'
@@ -465,7 +428,7 @@ def add_type_transition_rule(union_rule):
       # NB: This does not require that union bases be supplied to `def rules():`, as the union type
       # is never instantiated!
       union_base = union_rule.union_base
-      assert union_base._is_union
+      assert union.is_instance(union_base)
       union_member = union_rule.union_member
       if union_base not in union_rules:
         union_rules[union_base] = OrderedSet()

src/python/pants/engine/scheduler.py

@@ -22,7 +22,7 @@
 )
 from pants.engine.native import Function, TypeId
 from pants.engine.nodes import Return, Throw
-from pants.engine.objects import Collection
+from pants.engine.objects import Collection, union
 from pants.engine.rules import RuleIndex, TaskRule
 from pants.engine.selectors import Params
 from pants.util.contextutil import temporary_file_path
@@ -187,7 +187,7 @@ def add_get_edge(product, subject):
       self._native.lib.tasks_add_get(self._tasks, self._to_type(product), self._to_type(subject))
 
     for the_get in rule.input_gets:
-      if getattr(the_get.subject_declared_type, '_is_union', False):
+      if union.is_instance(the_get.subject_declared_type):
         # If the registered subject type is a union, add Get edges to all registered union members.
         for union_member in union_rules.get(the_get.subject_declared_type, []):
           add_get_edge(the_get.product, union_member)

src/python/pants/fs/fs.py

@@ -74,4 +74,3 @@ def expand_path(path):
 def is_child_of(path: Path, directory: Path) -> bool:
   abs_path = path if path.is_absolute() else directory.joinpath(path).resolve()
   return directory == abs_path or directory in abs_path.parents
-

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

@@ -9,7 +9,8 @@
 from pants.engine.fs import Digest, DirectoriesToMerge, DirectoryToMaterialize, Workspace
 from pants.engine.goal import Goal, GoalSubsystem, LineOriented
 from pants.engine.legacy.graph import HydratedTarget
-from pants.engine.rules import console_rule, rule, union
+from pants.engine.objects import union
+from pants.engine.rules import console_rule, rule
 from pants.engine.selectors import Get, MultiGet
 from pants.rules.core.distdir import DistDir
 

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

@@ -4,7 +4,7 @@
 from dataclasses import dataclass
 
 from pants.engine.isolated_process import FallibleExecuteProcessResult
-from pants.engine.rules import union
+from pants.engine.objects import union
 from pants.util.collections import Enum
 
 

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

@@ -10,7 +10,8 @@
 from pants.engine.isolated_process import ExecuteProcessResult
 from pants.engine.legacy.graph import HydratedTargets
 from pants.engine.legacy.structs import TargetAdaptor
-from pants.engine.rules import UnionMembership, console_rule, union
+from pants.engine.objects import union
+from pants.engine.rules import UnionMembership, console_rule
 from pants.engine.selectors import Get, MultiGet
 
 

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

@@ -8,8 +8,8 @@
 from pants.engine.isolated_process import FallibleExecuteProcessResult
 from pants.engine.legacy.graph import HydratedTargets
 from pants.engine.legacy.structs import TargetAdaptor
-from pants.engine.objects import Collection
-from pants.engine.rules import UnionMembership, console_rule, union
+from pants.engine.objects import Collection, union
+from pants.engine.rules import UnionMembership, console_rule
 from pants.engine.selectors import Get, MultiGet
 
 

src/python/pants/util/meta.py

@@ -1,6 +1,7 @@
 # Copyright 2015 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+from abc import ABC, abstractmethod
 from dataclasses import FrozenInstanceError
 from functools import wraps
 from typing import Any, Callable, Optional, Type, TypeVar, Union
@@ -122,6 +123,46 @@ def staticproperty(func: Callable[..., T]) -> T:
   return ClassPropertyDescriptor(func, doc)  # type: ignore[arg-type, return-value]
 
 
+class _ClassDecoratorWithSentinelAttribute(ABC):
+  """Base class to wrap a class decorator which sets a "sentinel attribute".
+
+  This functionality is exposed via the `@decorated_type_checkable` decorator.
+  """
+
+  @abstractmethod
+  def __call__(self, cls: Type) -> Type: ...
+
+  def define_instance_of(self, obj: Type, **kwargs) -> Type:
+    return type(obj.__name__, (obj,), {
+      '_decorated_type_checkable_type': type(self),
+      **kwargs
+    })
+
+  def is_instance(self, obj: Type) -> bool:
+    return getattr(obj, '_decorated_type_checkable_type', None) is type(self)
+
+
+def decorated_type_checkable(decorator: Callable[[Type], Type]) -> _ClassDecoratorWithSentinelAttribute:
+  """Wraps a class decorator to add a "sentinel attribute" to decorated classes.
+
+  A "sentinel attribute" is an attribute added to the wrapped class decorator's result with
+  `.define_instance_of()`. The wrapped class decorator can then be imported and used to check
+  whether some class object was wrapped with that decorator with `.is_instance()`.
+
+  When used on a class decorator method, the method should return
+  `<method name>.define_instance_of(cls)`, where `cls` is the class object that the decorator would
+  otherwise return.
+  """
+
+  class WrappedFunction(_ClassDecoratorWithSentinelAttribute):
+    @wraps(decorator)
+    def __call__(self, cls: Type) -> Type:
+      return decorator(cls)
+
+  return WrappedFunction()
+
+
+@decorated_type_checkable
 def frozen_after_init(cls: C) -> C:
   """Class decorator to freeze any modifications to the object after __init__() is done.
 
@@ -148,4 +189,5 @@ def new_setattr(self, key: str, value: Any) -> None:
 
   cls.__init__ = new_init
   cls.__setattr__ = new_setattr
+
   return cls

tests/python/pants_test/build_graph/test_build_configuration.py

@@ -10,7 +10,8 @@
 from pants.build_graph.build_configuration import BuildConfiguration
 from pants.build_graph.build_file_aliases import BuildFileAliases, TargetMacro
 from pants.build_graph.target import Target
-from pants.engine.rules import UnionRule, union
+from pants.engine.objects import union
+from pants.engine.rules import UnionRule
 from pants.util.contextutil import temporary_dir
 from pants.util.dirutil import touch
 

tests/python/pants_test/engine/BUILD

@@ -227,6 +227,7 @@ python_tests(
   sources=['test_objects.py'],
   dependencies=[
     'src/python/pants/engine:objects',
+    'src/python/pants/util:meta',
   ],
   tags = {"partially_type_checked"},
 )

tests/python/pants_test/engine/test_scheduler.py

@@ -10,7 +10,8 @@
 from typing import List
 
 from pants.engine.native import Native
-from pants.engine.rules import RootRule, UnionRule, rule, union
+from pants.engine.objects import union
+from pants.engine.rules import RootRule, UnionRule, rule
 from pants.engine.scheduler import ExecutionError, SchedulerSession
 from pants.engine.selectors import Get, Params
 from pants.testutil.engine.util import assert_equal_with_printing, remove_locations_from_traceback

tests/python/pants_test/util/test_meta.py

@@ -6,7 +6,13 @@
 from dataclasses import FrozenInstanceError, dataclass
 
 from pants.testutil.test_base import TestBase
-from pants.util.meta import SingletonMetaclass, classproperty, frozen_after_init, staticproperty
+from pants.util.meta import (
+  SingletonMetaclass,
+  classproperty,
+  decorated_type_checkable,
+  frozen_after_init,
+  staticproperty,
+)
 
 
 class AbstractClassTest(TestBase):
@@ -246,6 +252,34 @@ def f(cls):
     self.assertEqual(Concrete2.f, 'hello')
 
 
+class SentinelAttributeTest(unittest.TestCase):
+
+  def test_decorated_type_checkable(self):
+    @decorated_type_checkable
+    def f(cls):
+      return f.define_instance_of(cls)
+
+    @f
+    class C:
+      pass
+
+    self.assertEqual(C._decorated_type_checkable_type, type(f))
+    self.assertTrue(f.is_instance(C))
+
+    # Check that .is_instance() is only true for exactly the decorator @g used on the class D!
+    @decorated_type_checkable
+    def g(cls):
+      return g.define_instance_of(cls)
+
+    @g
+    class D:
+      pass
+
+    self.assertEqual(D._decorated_type_checkable_type, type(g))
+    self.assertTrue(g.is_instance(D))
+    self.assertFalse(f.is_instance(D))
+
+
 class FrozenAfterInitTest(unittest.TestCase):
 
   def test_no_init(self) -> None: