ensure @DataClass types have frozen=True when used in the engine

src/python/pants/engine/BUILD

@@ -147,6 +147,7 @@ python_library(
   sources=['rules.py'],
   dependencies=[
     '3rdparty/python:asttokens',
+    '3rdparty/python:dataclasses',
     '3rdparty/python/twitter/commons:twitter.common.collections',
     ':goal',
     ':selectors',

src/python/pants/engine/rules.py

@@ -9,21 +9,39 @@
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from collections.abc import Iterable
+from typing import Any, Dict, Set
 
 import asttokens
+import dataclasses
+from dataclasses import dataclass
 from twitter.common.collections import OrderedSet
 
 from pants.engine.goal import Goal
 from pants.engine.selectors import Get
 from pants.util.collections import assert_single_element
 from pants.util.memo import memoized
-from pants.util.objects import SubclassesOf, TypedCollection, datatype
+from pants.util.objects import TypeConstraint, TypeConstraintError, TypedCollection, datatype
 
 
 logger = logging.getLogger(__name__)
 
 
-_type_field = SubclassesOf(type)
+class _type_field_constraint(TypeConstraint):
+
+  def __init__(self):
+    super().__init__('<@dataclass `frozen=True` attribute-checking constraint>')
+
+  def satisfied_by(self, obj):
+    if not issubclass(type(obj), type):
+      return False
+    if dataclasses.is_dataclass(obj) and not obj.__dataclass_params__.frozen:
+      raise TypeConstraintError(
+        f'type {obj} is a dataclass declared without `frozen=True`! '
+        'The engine requires that fields in params are immutable for stable hashing!')
+    return True
+
+
+_type_field = _type_field_constraint()
 
 
 class _RuleVisitor(ast.NodeVisitor):
@@ -356,7 +374,7 @@ def dependency_optionables(self):
 
 class TaskRule(datatype([
   ('output_type', _type_field),
-  ('input_selectors', TypedCollection(SubclassesOf(type))),
+  ('input_selectors', TypedCollection(_type_field)),
   ('input_gets', tuple),
   'func',
   ('dependency_rules', tuple),
@@ -417,10 +435,14 @@ def dependency_optionables(self):
     return tuple()
 
 
-# TODO: add typechecking here -- would need to have a TypedCollection for dicts for `union_rules`.
-class RuleIndex(datatype(['rules', 'roots', 'union_rules'])):
+@dataclass(frozen=True)
+class RuleIndex:
   """Holds a normalized index of Rules used to instantiate Nodes."""
 
+  rules: Dict[Any, Any]
+  roots: Set[Any]
+  union_rules: Dict[Any, Any]
+
   @classmethod
   def create(cls, rule_entries, union_rules=None):
     """Creates a RuleIndex with tasks indexed by their output type."""

src/python/pants/util/objects.py

@@ -556,15 +556,9 @@ def exclude_iterable_constraint(cls):
   def __init__(self, constraint):
     """Create a `TypeConstraint` which validates each member of a collection with `constraint`.
 
-    :param TypeOnlyConstraint constraint: the `TypeConstraint` to apply to each element. This is
-                                          currently required to be a `TypeOnlyConstraint` to avoid
-                                          complex prototypal type relationships.
+    :param TypeConstraint constraint: the `TypeConstraint` to apply to each element.
     """
 
-    if not isinstance(constraint, TypeOnlyConstraint):
-      raise TypeError("constraint for collection must be a {}! was: {}"
-                      .format(TypeOnlyConstraint.__name__, constraint))
-
     description = '{}({})'.format(type(self).__name__, constraint)
 
     self._constraint = constraint

tests/python/pants_test/engine/test_scheduler.py

@@ -10,7 +10,7 @@
 from dataclasses import dataclass
 
 from pants.engine.native import Native
-from pants.engine.rules import RootRule, UnionRule, rule, union
+from pants.engine.rules import RootRule, RuleIndex, UnionRule, rule, union
 from pants.engine.scheduler import ExecutionError, SchedulerSession
 from pants.engine.selectors import Get, Params
 from pants.util.objects import datatype
@@ -59,6 +59,11 @@ def transitive_coroutine_rule(c):
   yield D(b)
 
 
+@dataclass
+class NonFrozenDataclass:
+  x: int
+
+
 @dataclass(frozen=True)
 class FrozenFieldsDataclass:
   x: int
@@ -376,3 +381,13 @@ def test_trace_includes_rule_exception_traceback(self):
                                raise Exception(f'An exception for {type(x).__name__}')
                            Exception: An exception for B''').lstrip() + '\n\n',  # Traces include two empty lines after.
                                trace)
+
+
+class RuleIndexingErrorTest(TestBase):
+
+  def test_non_frozen_dataclass_error(self):
+
+    with self.assertRaisesWithMessage(TypeError, dedent("""\
+      type check error in class RootRule: 1 error type checking constructor arguments:
+      field 'output_type' was invalid: type <class 'pants_test.engine.test_scheduler.NonFrozenDataclass'> is a dataclass declared without `frozen=True`! The engine requires that fields in params are immutable for stable hashing!""")):
+      RootRule(NonFrozenDataclass)

tests/python/pants_test/util/BUILD

@@ -104,7 +104,7 @@ python_tests(
 python_tests(
   name = 'objects-mypy',
   sources = ['test_objects_mypy.py'],
-  coverage = ['pants.util.objects.dataclasstype'],
+  coverage = ['pants.util.objects'],
   dependencies = [
     'testprojects/src/python/mypy:mypy-passing-examples',
     'testprojects/src/python/mypy:mypy-failing-examples',