Add @frozen_after_init decorator to make some @dataclass uses safer with V2

src/python/pants/util/BUILD

@@ -72,13 +72,16 @@ python_library(
   name = 'memo',
   sources = ['memo.py'],
   dependencies = [
-      ':meta'
+    ':meta'
   ],
 )
 
 python_library(
   name = 'meta',
   sources = ['meta.py'],
+  dependencies = [
+    '3rdparty/python:dataclasses'
+  ],
   tags = {'type_checked'},
 )
 

src/python/pants/util/meta.py

@@ -1,6 +1,8 @@
 # Copyright 2015 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+from dataclasses import FrozenInstanceError
+from functools import wraps
 from typing import Any, Callable, Optional, Type, TypeVar, Union
 
 
@@ -107,3 +109,32 @@ def staticproperty(func: Union[staticmethod, Callable]) -> ClassPropertyDescript
     func = staticmethod(func)
 
   return ClassPropertyDescriptor(func, doc)
+
+
+def freeze_after_init(cls: Type[T]) -> Type[T]:
+  """Class decorator to freeze any modifications to the object after __init__() is done.
+
+  The primary use case is for @dataclasses who cannot use frozen=True due to the need for a custom
+  __init__(), but who still want to remain as immutable as possible (e.g. for safety with the V2
+  engine). When using with dataclasses, this should be the first decorator applied, i.e. be used
+  before @dataclass."""
+
+  prev_init = cls.__init__
+  prev_setattr = cls.__setattr__
+
+  @wraps(prev_init)
+  def new_init(self, *args: Any, **kwargs: Any) -> None:
+    prev_init(self, *args, **kwargs)  # type: ignore
+    self._is_frozen = True
+
+  @wraps(prev_setattr)
+  def new_setattr(self, key: str, value: Any) -> None:
+    if getattr(self, "_is_frozen", False):
+      raise FrozenInstanceError(
+        f"Attempting to modify the attribute {key} after the object {self} was created."
+      )
+    prev_setattr(self, key, value)
+
+  cls.__init__ = new_init  # type: ignore
+  cls.__setattr__ = new_setattr  # type: ignore
+  return cls

tests/python/pants_test/util/BUILD

@@ -83,6 +83,7 @@ python_tests(
   name = 'meta',
   sources = ['test_meta.py'],
   dependencies = [
+    '3rdparty/python:dataclasses',
     'src/python/pants/util:meta',
     'tests/python/pants_test:test_base',
   ],

tests/python/pants_test/util/test_meta.py

@@ -1,9 +1,11 @@
 # Copyright 2015 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+import unittest
 from abc import ABC, abstractmethod
+from dataclasses import FrozenInstanceError, dataclass
 
-from pants.util.meta import SingletonMetaclass, classproperty, staticproperty
+from pants.util.meta import SingletonMetaclass, classproperty, freeze_after_init, staticproperty
 from pants_test.test_base import TestBase
 
 
@@ -242,3 +244,75 @@ class Concrete2(Abstract):
       def f(cls):
         return 'hello'
     self.assertEqual(Concrete2.f, 'hello')
+
+
+class FreezeAfterInitTest(unittest.TestCase):
+
+  def test_no_init(self) -> None:
+    @freeze_after_init
+    class Test:
+      pass
+
+    test = Test()
+    with self.assertRaises(FrozenInstanceError):
+      test.x = 1
+
+  def test_init_still_works(self):
+    @freeze_after_init
+    class Test:
+
+      def __init__(self, x: int) -> None:
+        self.x = x
+        self.y = "abc"
+
+    test = Test(x=0)
+    self.assertEqual(test.x, 0)
+    self.assertEqual(test.y, "abc")
+
+  def test_modify_preexisting_field_after_init(self) -> None:
+    @freeze_after_init
+    class Test:
+
+      def __init__(self, x: int) -> None:
+        self.x = x
+
+    test = Test(x=0)
+    with self.assertRaises(FrozenInstanceError):
+      test.x = 1
+
+  def test_add_new_field_after_init(self) -> None:
+    @freeze_after_init
+    class Test:
+
+      def __init__(self, x: int) -> None:
+        self.x: x
+
+    test = Test(x=0)
+    with self.assertRaises(FrozenInstanceError):
+      test.y = "abc"
+
+  def test_explicitly_call_setattr_after_init(self) -> None:
+    @freeze_after_init
+    class Test:
+
+      def __init__(self, x: int) -> None:
+        self.x: x
+
+    test = Test(x=0)
+    with self.assertRaises(FrozenInstanceError):
+      setattr(test, "x", 1)
+
+  def test_works_with_dataclass(self) -> None:
+    @freeze_after_init
+    @dataclass(frozen=False)
+    class Test:
+      x: int
+      y: str
+
+      def __init__(self, x: int) -> None:
+        self.x = x
+        self.y = "abc"
+
+    test = Test(x=0)
+    with self.assertRaises(FrozenInstanceError):
+      test.x = 1