serialization to json

src/drinx/__init__.py

@@ -3,6 +3,7 @@
 from drinx.attribute import field, static_field, private_field, static_private_field
 from drinx.visualize import visualize_leaf, tree_diagram, tree_summary
 from drinx.jax_utils import is_traced
+from drinx.serialize import to_dict, from_dict, to_json, from_json
 
 
 __all__ = [
@@ -16,4 +17,8 @@
     "tree_diagram",
     "tree_summary",
     "is_traced",
+    "to_dict",
+    "from_dict",
+    "to_json",
+    "from_json",
 ]

src/drinx/base.py

@@ -80,10 +80,14 @@ def __init_subclass__(
             unsafe_hash: Force generation of ``__hash__`` even when ``eq=True``.
             match_args: Set ``__match_args__`` for structural pattern matching.
             kw_only: Make all fields keyword-only in ``__init__``.
-            slots: Generate ``__slots__`` (ignored; kept for API compatibility).
+            slots: Not supported; ignored.  Slot-based dataclasses are excluded
+                because ``__slots__`` attribute-access speedups are negligible
+                compared to JAX kernel dispatch overhead, and supporting them
+                would add significant complexity to pytree flatten/unflatten.
             weakref_slot: Add a ``__weakref__`` slot (ignored; kept for API
                 compatibility).
         """
+        del slots, weakref_slot
         super().__init_subclass__()
         user_post_init = cls.__dict__.get("__post_init__")
         if user_post_init is not None and user_post_init is not DataClass.__post_init__:
@@ -102,7 +106,7 @@ def __init_subclass__(
             unsafe_hash=unsafe_hash,
             match_args=match_args,
             kw_only=kw_only,
-            weakref_slot=weakref_slot,
+            weakref_slot=False,
         )
         dataclass_transform(cls)
         setattr(cls, "__setattr__", DataClass.__setattr__)
@@ -374,11 +378,23 @@ def _traverse_path(
             attr_list.append(current_parent)
         return attr_list
 
+    @staticmethod
+    def _copy_and_set(obj: Any, field_name: str, value: Any) -> Any:
+        cls = type(obj)
+        new_obj = object.__new__(cls)
+        instance_dict = object.__getattribute__(obj, "__dict__")
+        for k, v in instance_dict.items():
+            object.__setattr__(new_obj, k, v)
+        object.__setattr__(new_obj, field_name, value)
+        return new_obj
+
     def aset(
         self,
         attr_name: str,
         val: Any,
         create_new_ok: bool = False,
+        allow_private: bool = False,
+        bypass_callbacks: bool = False,
     ) -> Self:
         """Sets an attribute of this class. In contrast to the classical .at[].set(), this method updates the class
         attribute directly and does not only operate on jax pytree leaf nodes. Instead, replaces the full attribute
@@ -397,6 +413,9 @@ def aset(
             val (Any): Value to set the attribute to
             create_new_ok (bool, optional): If false (default), throw an error if the attribute does not exist.
                 If true, creates a new attribute if the attribute name does not exist yet.
+            bypass_callbacks (bool, optional): If True, skip ``on_setattr`` callbacks for all attribute
+                operations in the path. If False (default), each attribute write runs the callbacks
+                registered on that field, mirroring the behaviour of ``__setattr__`` during ``__init__``.
 
         Returns:
             Self: Updated instance with new attribute value
@@ -410,11 +429,16 @@ def aset(
 
         # Validate the final step (respecting create_new_ok)
         if final_op_type == "attribute":
-            if not hasattr(parent, str(final_op)):
-                if not create_new_ok:
+            if dataclasses.is_dataclass(parent):
+                dc_field_names = {f.name for f in dataclasses.fields(parent)}
+                if str(final_op) not in dc_field_names and not create_new_ok:
                     raise Exception(
                         f"Attribute: {final_op} does not exist for {parent.__class__}"
                     )
+            elif not hasattr(parent, str(final_op)) and not create_new_ok:
+                raise Exception(
+                    f"Attribute: {final_op} does not exist for {parent.__class__}"
+                )
         elif final_op_type == "index":
             if not hasattr(parent, "__getitem__"):
                 raise Exception(f"{parent.__class__} does not implement __getitem__")
@@ -432,14 +456,28 @@ def aset(
             current_parent = attr_list[idx]
 
             if op_type == "attribute":
-                # Replaced generic DataClass check with standard dataclasses check
                 if not dataclasses.is_dataclass(current_parent):
                     raise Exception(
                         f"Can only set attribute functionally on a dataclass, but got {current_parent.__class__}"
                     )
 
-                # Use standard dataclasses.replace to functionally copy and update the frozen dataclass
-                cur_attr = dataclasses.replace(current_parent, **{str(op): cur_attr})
+                dc_fields = {f.name: f for f in dataclasses.fields(current_parent)}
+                target_field = dc_fields.get(str(op))
+                if target_field is None:
+                    raise TypeError(
+                        f"Field {str(op)!r} is not a dataclass field of {type(current_parent).__name__!r}."
+                    )
+                if not target_field.init and not allow_private:
+                    raise TypeError(
+                        f"Field {str(op)!r} has init=False (non-init/private field). "
+                        "Pass allow_private=True to allow updating non-init fields."
+                    )
+                if not bypass_callbacks:
+                    callbacks = DataClass._normalize_callbacks(
+                        target_field.metadata.get(DRINX_ON_SETATTR, ())
+                    )
+                    cur_attr = DataClass._run_callbacks(cur_attr, callbacks)
+                cur_attr = DataClass._copy_and_set(current_parent, str(op), cur_attr)
 
             elif op_type in ("index", "key"):
                 if not hasattr(current_parent, "copy"):
@@ -464,7 +502,13 @@ def aset(
         assert cur_attr.__class__ == self.__class__
         return cur_attr
 
-    def aset_inplace(self, attr_name: str, val: Any) -> None:
+    def aset_inplace(
+        self,
+        attr_name: str,
+        val: Any,
+        create_new_ok: bool = False,
+        bypass_callbacks: bool = False,
+    ) -> None:
         """Sets an attribute of this dataclass **in place** by bypassing the frozen
         restriction via ``object.__setattr__``.
 
@@ -485,12 +529,46 @@ def aset_inplace(self, attr_name: str, val: Any) -> None:
         Args:
             attr_name: Path string (see :meth:`aset` for syntax).
             val: Value to assign at the target location.
+            create_new_ok (bool, optional): If False (default), raise if the target
+                attribute or dictionary key does not already exist.  If True, allow
+                setting new attributes or creating new dictionary keys.
+            bypass_callbacks (bool, optional): If True, skip ``on_setattr``
+                callbacks. If False, fire the callbacks registered on the target field before
+                writing, mirroring the behaviour of :meth:`aset` (default).
         """
         ops = self._parse_operations(attr_name)
         attr_list = self._traverse_path(ops)
         final_op, final_op_type = ops[-1]
         parent = attr_list[-1]
 
+        if final_op_type == "attribute":
+            if dataclasses.is_dataclass(parent):
+                dc_field_names = {f.name for f in dataclasses.fields(parent)}
+                if str(final_op) not in dc_field_names and not create_new_ok:
+                    raise Exception(
+                        f"Attribute: {final_op} does not exist for {parent.__class__}"
+                    )
+            elif not hasattr(parent, str(final_op)) and not create_new_ok:
+                raise Exception(
+                    f"Attribute: {final_op} does not exist for {parent.__class__}"
+                )
+        elif final_op_type == "key":
+            if not hasattr(parent, "__getitem__"):
+                raise Exception(f"{parent.__class__} does not implement __getitem__")
+            if final_op not in parent:
+                if not create_new_ok:
+                    raise Exception(f"Key: {final_op} does not exist for {parent}")
+
+        if final_op_type == "attribute" and not bypass_callbacks:
+            if dataclasses.is_dataclass(parent):
+                dc_fields = {f.name: f for f in dataclasses.fields(parent)}
+                target_field = dc_fields.get(str(final_op))
+                if target_field is not None:
+                    callbacks = DataClass._normalize_callbacks(
+                        target_field.metadata.get(DRINX_ON_SETATTR, ())
+                    )
+                    val = DataClass._run_callbacks(val, callbacks)
+
         if final_op_type == "attribute":
             object.__setattr__(parent, str(final_op), val)
         elif final_op_type == "index":
@@ -537,7 +615,9 @@ def __init__(self, obj: _DC, path: list[str | int | Any]) -> None:
     def __getitem__(self, key: Any) -> "_AtIndexer[_DC]":
         return _AtIndexer(self._obj, self._path + [key])
 
-    def set(self, value: Any) -> _DC:
+    def set(
+        self, value: Any, allow_private: bool = False, bypass_callbacks: bool = False
+    ) -> _DC:
         """Apply a functional update and return the new DataClass instance.
 
         If the accumulated path contains only ``str``/``int`` keys, delegates to
@@ -547,6 +627,7 @@ def set(self, value: Any) -> _DC:
         Args:
             value: The new value.  For mask updates this may be a scalar or a
                 DataClass tree of the same type as the mask/object.
+            bypass_callbacks: If True, skip ``on_setattr`` callbacks. Default False runs them.
 
         Returns:
             Updated DataClass instance.
@@ -585,7 +666,12 @@ def set(self, value: Any) -> _DC:
                     "Expected str, int, or a DataClass mask."
                 )
         path_str = "->".join(parts)
-        return self._obj.aset(path_str, value)
+        return self._obj.aset(
+            path_str,
+            value,
+            allow_private=allow_private,
+            bypass_callbacks=bypass_callbacks,
+        )
 
 
 class _AtProxy(Generic[_DC]):