feat(py_class): support frozen=True for immutable instances (#542)

## Summary

- Add `frozen` parameter to `@py_class()` decorator and `field()`
function, mirroring Python's `dataclasses.dataclass(frozen=True)`
semantics.
- **Class-level frozen**: `@py_class(frozen=True)` installs
`__setattr__`/`__delattr__` guards raising `FrozenInstanceError`. Frozen
classes auto-get `__hash__` (safely hashable).
- **Field-level frozen**: `field(frozen=True)` sets
`TypeField.frozen=True` (`fset=None`), independent of class-level
frozen.
- `__replace__` (`copy.replace`) uses direct `FieldSetter` for
`py_class` (bypasses both frozen mechanisms) while `c_class` respects
C++ readonly fields via `object.__setattr__`.

## Test plan

- [x] 101 new tests in `tests/python/test_dataclass_frozen.py` covering:
  - Class-level frozen basic behavior and error messages
  - Field-level frozen (individual field override)
  - Inheritance (frozen parent + mutable child, mixed)
  - Interactions with `eq`, `hash`, `order`, `copy`, `replace`
  - `object.__setattr__` escape hatch
  - `FrozenInstanceError` is `AttributeError` subclass
- [x] Full test suite: 2105 passed, 0 failed, 38 skipped, 3 xfailed

Author: junrushao

python/tvm_ffi/_dunder.py

@@ -193,8 +193,9 @@ def _make_replace(_type_info: TypeInfo) -> Callable[..., Any]:
 
     def __replace__(self: Any, **kwargs: Any) -> Any:
         obj = copy_copy(self)
+        cls = type(obj)
         for key, value in kwargs.items():
-            setattr(obj, key, value)
+            getattr(cls, key).set(obj, value)
         return obj
 
     return __replace__

python/tvm_ffi/cython/type_info.pxi

@@ -593,6 +593,24 @@ def _annotation_cobject(cls, targs):
     return TypeSchema(origin, origin_type_index=info.type_index)
 
 
+class FFIProperty(property):
+    """Property descriptor for FFI-backed fields.
+
+    When *frozen* is True the public setter (``fset``) is suppressed so
+    that normal attribute assignment raises ``AttributeError``.  The
+    real setter is stashed in :attr:`_fset` and exposed via the
+    :meth:`set` escape-hatch.
+    """
+
+    def __init__(self, fget, fset, frozen, fdel=None, doc=None):
+        super().__init__(fget, None if frozen else fset, fdel, doc)
+        self._fset = fset
+
+    def set(self, obj, value):
+        """Force-set the field value, bypassing the frozen guard."""
+        self._fset(obj, value)
+
+
 @dataclasses.dataclass(eq=False)
 class TypeField:
     """Description of a single reflected field on an FFI-backed type."""
@@ -616,17 +634,18 @@ class TypeField:
         assert self.getter is not None
 
     def as_property(self, object cls):
-        """Create a Python ``property`` object for this field on ``cls``."""
+        """Create an :class:`FFIProperty` descriptor for this field on ``cls``."""
         cdef str name = self.name
         cdef FieldGetter fget = self.getter
         cdef FieldSetter fset = self.setter
         cdef object ret
         fget.__name__ = fset.__name__ = name
         fget.__module__ = fset.__module__ = cls.__module__
         fget.__qualname__ = fset.__qualname__ = f"{cls.__qualname__}.{name}"
-        ret = property(
+        ret = FFIProperty(
             fget=fget,
-            fset=fset if (not self.frozen) else None,
+            fset=fset,
+            frozen=self.frozen,
         )
         if self.doc:
             ret.__doc__ = self.doc
@@ -1003,7 +1022,7 @@ def _register_fields(type_info, fields, structure_kind=None):
                 doc=py_field.doc,
                 size=size,
                 offset=field_offset,
-                frozen=False,
+                frozen=py_field.frozen,
                 metadata={"type_schema": py_field.ty.to_json()},
                 getter=fgetter,
                 setter=fsetter,

python/tvm_ffi/dataclasses/field.py

@@ -58,6 +58,8 @@ class Field:
     default_factory : Callable[[], object] | None
         A zero-argument callable that produces the default value.
         Mutually exclusive with *default*.  ``None`` when not set.
+    frozen : bool
+        Whether this field is read-only after ``__init__``.
     init : bool
         Whether this field appears in the auto-generated ``__init__``.
     repr : bool
@@ -91,6 +93,7 @@ class Field:
         "default",
         "default_factory",
         "doc",
+        "frozen",
         "hash",
         "init",
         "kw_only",
@@ -103,6 +106,7 @@ class Field:
     ty: TypeSchema | None
     default: object
     default_factory: Callable[[], object] | None
+    frozen: bool
     init: bool
     repr: bool
     hash: bool | None
@@ -123,6 +127,7 @@ def __init__(  # noqa: PLR0913
         *,
         default: object = MISSING,
         default_factory: Callable[[], object] | None = MISSING,  # type: ignore[assignment]
+        frozen: bool = False,
         init: bool = True,
         repr: bool = True,
         hash: bool | None = True,
@@ -151,6 +156,7 @@ def __init__(  # noqa: PLR0913
         self.ty = ty
         self.default = default
         self.default_factory = default_factory
+        self.frozen = frozen
         self.init = init
         self.repr = repr
         self.hash = hash
@@ -164,6 +170,7 @@ def field(
     *,
     default: object = MISSING,
     default_factory: Callable[[], object] | None = MISSING,  # type: ignore[assignment]
+    frozen: bool = False,
     init: bool = True,
     repr: bool = True,
     hash: bool | None = None,
@@ -189,6 +196,11 @@ def field(
     default_factory
         A zero-argument callable that produces the default value.
         Mutually exclusive with *default*.
+    frozen
+        Whether this field is read-only after ``__init__``.  When True,
+        the Python property descriptor has no setter; use the
+        ``type(obj).field_name.set(obj, value)`` escape hatch when
+        mutation is necessary.
     init
         Whether this field appears in the auto-generated ``__init__``.
     repr
@@ -234,6 +246,7 @@ class MyFunc(Object):
     return Field(
         default=default,
         default_factory=default_factory,
+        frozen=frozen,
         init=init,
         repr=repr,
         hash=hash,

python/tvm_ffi/dataclasses/py_class.py

@@ -134,10 +134,11 @@ def _rollback_registration(cls: type, type_info: Any) -> None:
 # ---------------------------------------------------------------------------
 
 
-def _collect_own_fields(
+def _collect_own_fields(  # noqa: PLR0912
     cls: type,
     hints: dict[str, Any],
     decorator_kw_only: bool,
+    decorator_frozen: bool,
 ) -> list[Field]:
     """Parse own annotations into :class:`Field` objects.
 
@@ -194,6 +195,10 @@ def _collect_own_fields(
         if f.kw_only is None:
             f.kw_only = kw_only_active
 
+        # Apply class-level frozen when the field doesn't explicitly set it
+        if decorator_frozen and not f.frozen:
+            f.frozen = True
+
         # Resolve hash=None → follow compare (native dataclass semantics)
         if f.hash is None:
             f.hash = f.compare
@@ -248,7 +253,7 @@ def _register_fields_into_type(
     except (NameError, AttributeError):
         return False
 
-    own_fields = _collect_own_fields(cls, hints, params["kw_only"])
+    own_fields = _collect_own_fields(cls, hints, params["kw_only"], params["frozen"])
     py_methods = _collect_py_methods(cls)
 
     # Register fields and type-level structural eq/hash kind with the C layer.
@@ -414,11 +419,12 @@ def _install_deferred_init(
     order_default=False,
     field_specifiers=(field, Field),
 )
-def py_class(
+def py_class(  # noqa: PLR0913
     cls_or_type_key: type | str | None = None,
     /,
     *,
     type_key: str | None = None,
+    frozen: bool = False,
     init: bool = True,
     repr: bool = True,
     eq: bool = False,
@@ -465,6 +471,11 @@ class MyNode(Object):
     type_key
         Explicit FFI type key.  Auto-generated from
         ``{module}.{qualname}`` when omitted.
+    frozen
+        If True, all fields are read-only after ``__init__`` by default.
+        Individual fields can still be marked ``field(frozen=True)`` on a
+        non-frozen class.  Use ``type(obj).field_name.set(obj, value)``
+        as an escape hatch when mutation is necessary.
     init
         If True (default), generate ``__init__`` from field annotations.
     repr
@@ -514,6 +525,7 @@ class MyNode(Object):
 
     effective_type_key = type_key
     params: dict[str, Any] = {
+        "frozen": frozen,
         "init": init,
         "repr": repr,
         "eq": eq,

tests/python/test_cubin_launcher.py

@@ -93,7 +93,7 @@ def _compile_kernel_to_cubin() -> bytes:
         )
 
         if result.returncode != 0:
-            pytest.skip(f"nvcc not available or compilation failed: {result.stderr}")
+            pytest.skip(f"nvcc not available or compilation failed: {result.stderr}")  # ty: ignore[invalid-argument-type, too-many-positional-arguments]
 
         return cubin_file.read_bytes()
 

tests/python/test_dataclass_copy.py

@@ -934,10 +934,14 @@ def test_original_unchanged(self) -> None:
         obj.__replace__(v_i64=100)  # ty: ignore[unresolved-attribute]
         assert obj.v_i64 == 5  # ty: ignore[unresolved-attribute]
 
-    def test_replace_readonly_field_raises(self) -> None:
+    def test_replace_readonly_field(self) -> None:
+        # __replace__ uses the FFIProperty.set() escape hatch,
+        # so it works even on frozen / read-only fields.
         pair = tvm_ffi.testing.TestIntPair(3, 4)
-        with pytest.raises(AttributeError):
-            pair.__replace__(a=10)  # ty: ignore[unresolved-attribute]
+        pair2 = pair.__replace__(a=10)  # ty: ignore[unresolved-attribute]
+        assert pair2.a == 10
+        assert pair2.b == 4
+        assert pair.a == 3  # original unchanged
 
     def test_auto_replace_for_cxx_class(self) -> None:
         # _TestCxxClassBase is copy-constructible, so replace is auto-enabled