feat(frozen_dataclass): implement enhanced frozen_dataclass with post-init mutability and explicit sealing

Author: tony

src/libtmux/_internal/frozen_dataclass.py

@@ -10,147 +10,281 @@
 import dataclasses
 import functools
 import typing as t
+from dataclasses import field
 
 from typing_extensions import dataclass_transform
 
 _T = t.TypeVar("_T")
 
 
+def mutable_during_init(
+    default: t.Any = None, *, default_factory: t.Callable[[], t.Any] | None = None
+) -> dataclasses.Field[t.Any]:
+    """Create a field that can be modified during initialization phase.
+
+    This field can be modified after __init__ but before the object is sealed.
+    This is useful for establishing circular references between objects.
+
+    Parameters
+    ----------
+    default : Any, optional
+        Default value for the field
+    default_factory : callable, optional
+        A function that returns the default value
+
+    Returns
+    -------
+    dataclasses.Field
+        A field with metadata indicating it's mutable during initialization
+
+    Examples
+    --------
+    >>> @frozen_dataclass
+    ... class Node:
+    ...     name: str
+    ...     parent: Node | None = mutable_during_init(None)
+    ...     children: list[Node] = mutable_during_init(default_factory=list)
+    """
+    metadata = {"mutable_during_init": True}
+    result: dataclasses.Field[t.Any]
+    if default_factory is not None:
+        result = field(default_factory=default_factory, metadata=metadata)
+    else:
+        result = field(default=default, metadata=metadata)
+    return result
+
+
 @dataclass_transform(frozen_default=True)
-def frozen_dataclass(cls: type[_T]) -> type[_T]:
-    """Create a dataclass that's effectively immutable but inherits from non-frozen.
+def frozen_dataclass(cls: type[_T] | None = None, *, auto_seal: bool = False) -> t.Any:
+    """Create a dataclass that's effectively immutable but with controlled mutability.
 
     This decorator:
       1) Applies dataclasses.dataclass(frozen=False) to preserve normal dataclass
          generation
-      2) Overrides __setattr__ and __delattr__ to block changes post-init
-      3) Tells type-checkers that the resulting class should be treated as frozen
+      2) Identifies fields marked as mutable_during_init via field metadata
+      3) Allows controlled mutation of those fields after __init__ until sealed
+      4) Provides a seal() method to make the object completely immutable
+      5) Tells type-checkers that the resulting class should be treated as frozen
 
     Parameters
     ----------
-    cls : Type[_T]
-        The class to convert to a frozen-like dataclass
+    cls : Type[_T], optional
+        The class to convert to a controlled-frozen dataclass
+    auto_seal : bool, default=False
+        If True, automatically seal the object after __post_init__
+        If False, the user must call seal() explicitly when done with initialization
 
     Returns
     -------
     Type[_T]
-        The processed class with immutability enforced at runtime
+        The processed class with controlled immutability enforced at runtime
 
     Examples
     --------
-    Basic usage:
+    Basic usage with explicit sealing:
 
     >>> @frozen_dataclass
     ... class User:
     ...     id: int
     ...     name: str
+    ...     friends: list[str] = mutable_during_init(default_factory=list)
     >>> user = User(id=1, name="Alice")
-    >>> user.name
-    'Alice'
-    >>> user.name = "Bob"
+    >>> user.friends.append("Bob")  # Modifying the list contents (always allowed)
+    >>> user.friends = ["Bob", "Charlie"]  # Allowed for mutable_during_init fields
+    >>> user.name = "Alice Smith"  # Raises error - normal fields can't be modified
     Traceback (most recent call last):
         ...
     AttributeError: User is immutable: cannot modify field 'name'
+    >>> user.seal()  # Finalize the object, making it completely immutable
+    >>> user.friends = ["Dave"]  # Now raises error after sealing
+    Traceback (most recent call last):
+        ...
+    AttributeError: User is sealed: cannot modify any fields
+
+    Circular references with explicit sealing:
 
-    Mutating internal attributes (_-prefixed):
+    >>> @frozen_dataclass
+    ... class Node:
+    ...     name: str
+    ...     parent: 'Node | None' = mutable_during_init(None)
+    ...     children: list['Node'] = mutable_during_init(default_factory=list)
+    >>> root = Node(name="root")
+    >>> child1 = Node(name="child1")
+    >>> child2 = Node(name="child2")
+    >>> # Establish the parent-child relationships
+    >>> child1.parent = root
+    >>> child2.parent = root
+    >>> root.children.append(child1)
+    >>> root.children.append(child2)
+    >>> # Once relationships are established, seal all nodes
+    >>> root.seal()
+    >>> child1.seal()
+    >>> child2.seal()
+    >>> # Now the structure is immutable
+    >>> child3 = Node(name="child3")
+    >>> root.children.append(child3)  # List contents are still mutable
+    >>> child3.parent = root  # OK - child3 is not sealed yet
+    >>> try:
+    ...     root.name = "ROOT"  # Error - root is sealed
+    ... except AttributeError:
+    ...     print("Cannot modify sealed object")
+    Cannot modify sealed object
 
-    >>> user._cache = {"logged_in": True}
-    >>> user._cache
-    {'logged_in': True}
+    Auto-sealing after initialization:
 
-    Nested mutable fields limitation:
+    >>> @frozen_dataclass(auto_seal=True)
+    ... class Config:
+    ...     values: dict[str, str] = field(default_factory=dict)
+    ...     def __post_init__(self):
+    ...         self.values["timestamp"] = "now"
+    ...         # Object will be auto-sealed after __post_init__
+    >>> conf = Config()
+    >>> conf.values["key"] = "value"  # Dict contents are still mutable
+    >>> try:
+    ...     conf.values = {}  # Error - object was auto-sealed
+    ... except AttributeError:
+    ...     print("Cannot modify sealed object")
+    Cannot modify sealed object
+
+    Internal attributes are always mutable:
 
     >>> @frozen_dataclass
-    ... class Container:
-    ...     items: list[int]
-    >>> c = Container(items=[1, 2])
-    >>> c.items.append(3)  # allowed; mutable field itself isn't protected
-    >>> c.items
-    [1, 2, 3]
-    >>> # For deep immutability, use immutable collections (tuple, frozenset)
-    >>> @frozen_dataclass
-    ... class ImmutableContainer:
-    ...     items: tuple[int, ...] = (1, 2)
-    >>> ic = ImmutableContainer()
-    >>> ic.items
-    (1, 2)
-
-    Inheritance from mutable base classes:
-
-    >>> import dataclasses
-    >>> @dataclasses.dataclass
-    ... class MutableBase:
-    ...     value: int
-    >>> @frozen_dataclass
-    ... class ImmutableSub(MutableBase):
-    ...     pass
-    >>> obj = ImmutableSub(42)
-    >>> obj.value
-    42
-    >>> obj.value = 100
-    Traceback (most recent call last):
-        ...
-    AttributeError: ImmutableSub is immutable: cannot modify field 'value'
+    ... class Cache:
+    ...     name: str
+    >>> cache = Cache(name="mycache")
+    >>> cache.seal()
+    >>> cache._internal = {"data": 123}  # Private attributes are allowed
+    >>> cache._internal
+    {'data': 123}
 
-    Security consideration - modifying the _frozen flag:
+    CAUTION: The implementation allows setting _sealed=False to bypass immutability:
 
     >>> @frozen_dataclass
     ... class SecureData:
     ...     secret: str
     >>> data = SecureData(secret="password123")
-    >>> data.secret = "hacked"
-    Traceback (most recent call last):
-        ...
-    AttributeError: SecureData is immutable: cannot modify field 'secret'
-    >>> # CAUTION: The _frozen attribute can be modified to bypass immutability
-    >>> # protection. This is a known limitation of this implementation
-    >>> data._frozen = False  # intentionally bypassing immutability
-    >>> data.secret = "hacked"  # now works because object is no longer frozen
+    >>> data.seal()
+    >>> # CAUTION: This bypasses the immutability protection
+    >>> data._frozen = False  # Intentionally bypassing immutability
+    >>> data.secret = "hacked"  # Works because object is no longer sealed
     >>> data.secret
     'hacked'
     """
-    # A. Convert to a dataclass with frozen=False
-    cls = dataclasses.dataclass(cls)
-
-    # B. Explicitly annotate and initialize the `_frozen` attribute for static analysis
-    cls.__annotations__["_frozen"] = bool
-    setattr(cls, "_frozen", False)
-
-    # Save the original __init__ to use in our hooks
-    original_init = cls.__init__
-
-    # C. Create a new __init__ that will call the original and then set _frozen flag
-    @functools.wraps(original_init)
-    def __init__(self: t.Any, *args: t.Any, **kwargs: t.Any) -> None:
-        # Call the original __init__
-        original_init(self, *args, **kwargs)
-        # Set the _frozen flag to make object immutable
-        object.__setattr__(self, "_frozen", True)
-
-    # D. Custom attribute assignment method
-    def __setattr__(self: t.Any, name: str, value: t.Any) -> None:
-        # If _frozen is set and we're trying to set a field, block it
-        if getattr(self, "_frozen", False) and not name.startswith("_"):
-            # Allow mutation of private (_-prefixed) attributes after initialization
+
+    def wrap(cls: type[_T]) -> type[_T]:
+        # A. Convert to a dataclass with frozen=False
+        cls = dataclasses.dataclass(cls)
+
+        # B. Collect fields marked as mutable during initialization
+        mutable_during_init_fields = set()
+        for name, field_obj in getattr(cls, "__dataclass_fields__", {}).items():
+            if field_obj.metadata.get("mutable_during_init", False):
+                mutable_during_init_fields.add(name)
+
+        # Store on class for runtime checks
+        setattr(cls, "_mutable_during_init_fields", mutable_during_init_fields)
+
+        # C. Add state tracking attributes
+        cls.__annotations__["_initializing"] = bool
+        cls.__annotations__["_sealed"] = bool
+        cls.__annotations__["_frozen"] = bool  # For backward compatibility
+        setattr(cls, "_initializing", True)
+        setattr(cls, "_sealed", False)
+        setattr(cls, "_frozen", True)  # Backward compatibility
+
+        # Save original methods
+        original_init = cls.__init__
+        original_post_init = getattr(cls, "__post_init__", None)
+
+        # D. Custom __init__ to set initial state
+        @functools.wraps(original_init)
+        def __init__(self: t.Any, *args: t.Any, **kwargs: t.Any) -> None:
+            # Set initializing state
+            object.__setattr__(self, "_initializing", True)
+            object.__setattr__(self, "_sealed", False)
+            object.__setattr__(self, "_frozen", False)
+
+            # Call original __init__
+            original_init(self, *args, **kwargs)
+
+            # After __init__, only mutable_during_init fields can be changed
+            object.__setattr__(self, "_initializing", False)
+            object.__setattr__(self, "_frozen", True)
+
+            # If there's a post_init, it will be called separately
+
+        # E. Custom __post_init__ to potentially auto-seal
+        def __post_init__(self: t.Any) -> None:
+            if original_post_init:
+                original_post_init(self)
+
+            if auto_seal:
+                self.seal()
+
+        # F. Add seal method
+        def seal(self: t.Any) -> None:
+            """Finalize this object, making it completely immutable."""
+            object.__setattr__(self, "_sealed", True)
+
+        # G. Custom attribute setting to enforce rules
+        def __setattr__(self: t.Any, name: str, value: t.Any) -> None:
+            # Allow setting during initialization phase
+            if getattr(self, "_initializing", True):
+                object.__setattr__(self, name, value)
+                return
+
+            # Backward compatibility: if _frozen was explicitly set to False,
+            # allow mutations
+            if (
+                name not in ("_frozen", "_sealed")
+                and getattr(self, "_frozen", True) is False
+            ):
+                object.__setattr__(self, name, value)
+                return
+
+            # Always allow internal attributes
+            if name.startswith("_"):
+                object.__setattr__(self, name, value)
+                return
+
+            # If sealed, no attributes can be modified
+            if getattr(self, "_sealed", False):
+                error_msg = f"{cls.__name__} is sealed: cannot modify any fields"
+                raise AttributeError(error_msg)
+
+            # If not sealed but past init, only specially marked fields can be changed
+            if name in getattr(self, "_mutable_during_init_fields", set()):
+                object.__setattr__(self, name, value)
+                return
+
+            # Otherwise, field cannot be modified
             error_msg = f"{cls.__name__} is immutable: cannot modify field '{name}'"
             raise AttributeError(error_msg)
 
-        # Allow the assignment
-        object.__setattr__(self, name, value)
+        # H. Custom attribute deletion to prevent deletion
+        def __delattr__(self: t.Any, name: str) -> None:
+            if name.startswith("_") and not getattr(self, "_sealed", False):
+                object.__delattr__(self, name)
+                return
+
+            if getattr(self, "_sealed", False):
+                error_msg = f"{cls.__name__} is sealed: cannot delete any fields"
+                raise AttributeError(error_msg)
 
-    # E. Custom attribute deletion method
-    def __delattr__(self: t.Any, name: str) -> None:
-        # If we're frozen, block deletion
-        if getattr(self, "_frozen", False):
             error_msg = f"{cls.__name__} is immutable: cannot delete field '{name}'"
             raise AttributeError(error_msg)
 
-        # Allow the deletion
-        object.__delattr__(self, name)
+        # I. Inject methods
+        setattr(cls, "__init__", __init__)
+        setattr(cls, "__post_init__", __post_init__)
+        setattr(cls, "seal", seal)
+        setattr(cls, "__setattr__", __setattr__)
+        setattr(cls, "__delattr__", __delattr__)
+
+        return cls
 
-    # F. Inject methods into the class (using setattr to satisfy mypy)
-    setattr(cls, "__init__", __init__)  # Sets _frozen flag post-initialization
-    setattr(cls, "__setattr__", __setattr__)  # Blocks attribute modification post-init
-    setattr(cls, "__delattr__", __delattr__)  # Blocks attribute deletion post-init
+    # Support both @frozen_dataclass and @frozen_dataclass(...) syntax
+    if cls is None:
+        return wrap
 
-    return cls
+    return wrap(cls)