ENH support subclasses of list, set, dict, object for circular pointers

docs/changes.rst

@@ -26,17 +26,33 @@ v0.17
   ``Series`` or ``DataFrame`` using one, could not be saved. Such objects are
   now saved through ``__getstate__``/``__dict__`` again, as before v0.12.0.
   :pr:`550` by `Adrin Jalali`_.
-- Objects which contain a reference to themselves, directly or through
-  their attributes, dicts, lists, or sets, can now be saved and loaded. They
-  used to fail with a ``RecursionError``; this affected for instance the
-  discrete distributions of ``scipy.stats`` and fitted
-  :class:`sklearn.cluster.Birch` models, which is no longer listed as
-  unsupported. Such a reference through any other type, e.g. a tuple, raises
-  an ``UnsupportedTypeException`` when saving. Bound methods inherited from a
-  class defined in another module can now be loaded, they used to be
-  rejected as corrupted. A file in which nodes of different types share an
-  ``__id__`` is now rejected when loading instead of silently loading one of
-  them in place of the other. :pr:`549` by `Adrin Jalali`_.
+- Objects which contain a reference to themselves can now be saved and
+  loaded, and subclasses of the built-in containers are persisted the way
+  pickle does it. :pr:`549` and :pr:`554` by `Adrin Jalali`_.
+
+  - Circular references: objects which refer to themselves, directly or
+    through their attributes, dicts, lists, or sets, including subclasses of
+    these, can now be saved and loaded. They used to fail with a
+    ``RecursionError``; this affected for instance the discrete distributions
+    of ``scipy.stats`` and fitted :class:`sklearn.cluster.Birch` models, which
+    is no longer listed as unsupported. Such a reference through any other
+    type, e.g. a tuple, raises an ``UnsupportedTypeException`` when saving.
+
+  - Container subclasses: subclasses of ``dict``, ``list`` and ``set`` are now
+    saved and loaded the way pickle does it for dicts and lists. The instance
+    is created with ``__new__`` and filled in place instead of through its
+    constructor, so a subclass whose constructor requires arguments can now
+    be loaded, and its instance attributes, which used to be dropped, are
+    saved and restored. Subclasses of ``set`` and of
+    ``collections.defaultdict`` are treated the same way, although pickle
+    builds them through their constructor, so that items which refer back to
+    the container can be loaded. Files written with an earlier protocol load
+    as before.
+
+  - Loading: bound methods inherited from a class defined in another module
+    can now be loaded, they used to be rejected as corrupted. A file in which
+    nodes of different types share an ``__id__`` is now rejected instead of
+    silently loading one of them in place of the other.
 - Restore the ``skops`` command line entry point. It was declared in
   ``setup.py`` and lost when the packaging moved to ``pyproject.toml`` in
   v0.11.0, so ``skops convert`` and ``skops update`` had not been available
@@ -48,6 +64,12 @@ v0.17
   could not be loaded. The file format now stores the keyword arguments and
   the persistence protocol is bumped to 3; files written with an earlier
   protocol load as before. :pr:`551` by `Adrin Jalali`_.
+- Fix loading of ``collections.defaultdict`` objects whose keys are not
+  strings, which failed with a ``TypeError``, and of objects whose class
+  defines ``__slots__``, which failed with an ``AttributeError`` on Python
+  3.11 and later. Subclasses of ``defaultdict`` are now loaded as their own
+  type, with their instance attributes; they used to be loaded as a plain
+  ``defaultdict``. :pr:`554` by `Adrin Jalali`_.
 
 v0.16
 -----

skops/io/_audit.py

@@ -173,8 +173,8 @@ def construct(self) -> Any:
 
         A node is reached again while its own ``_construct`` runs when the saved
         object contained a reference to itself, directly or through its
-        children. ``ObjectNode``, ``DictNode``, ``ListNode`` and ``SetNode``
-        support this by storing the instance in ``_constructed`` before
+        children. ``ObjectNode``, ``DictNode``, ``DefaultDictNode``, ``ListNode``
+        and ``SetNode`` support this by storing the instance in ``_constructed`` before
         constructing the children, so the second call returns the partially
         constructed instance, like pickle does. Any other node raises instead
         of recursing until the interpreter gives up.

skops/io/_general.py

@@ -40,6 +40,73 @@
 arepr.maxstring = 24
 
 
+def _get_attrs_state(obj: Any, save_context: SaveContext) -> dict[str, Any] | None:
+    """Return the state of the instance attributes of a container, if any.
+
+    A plain dict, list or set has no instance attributes. An instance of a
+    subclass has a ``__dict__``, unless the class defines ``__slots__``, and
+    its attributes are read the way ``object_get_state`` reads them: through
+    ``__getstate__`` when the object has one, and from ``__dict__`` otherwise.
+    Every object has a ``__getstate__`` since Python 3.11, and the default one
+    returns ``None`` when there is nothing to save. A class which defines its
+    own decides what is worth saving, an empty dict included, and gets it back
+    through ``__setstate__``, as with pickle. ``None`` is returned when there
+    is nothing to save, so that the state of an object without attributes is
+    the same on every Python version.
+    """
+    if hasattr(obj, "__getstate__"):
+        attrs = obj.__getstate__()
+    else:
+        # Python < 3.11, where only a custom ``__getstate__`` exists: an empty
+        # ``__dict__`` is what the default ``__getstate__`` reports as ``None``.
+        attrs = getattr(obj, "__dict__", None) or None
+    if attrs is None:
+        return None
+    return get_state(attrs, save_context)
+
+
+def _get_attrs_tree(
+    state: dict[str, Any], load_context: LoadContext, trusted: TrustedTypes | None
+) -> Node | None:
+    """Return the node of the ``attrs`` entry of ``state``, if there is one.
+
+    The entry is written by ``_get_attrs_state``, for instances which have
+    attributes.
+    """
+    attrs = state.get("attrs")
+    if attrs is None:
+        return None
+    return get_tree(attrs, load_context, trusted=trusted)
+
+
+def _set_attrs(instance: Any, attrs: Node | None) -> None:
+    """Give ``instance`` its attributes back from the node of their state.
+
+    ``ObjectNode``, ``DictNode``, ``ListNode`` and ``SetNode`` create the
+    instance with ``__new__`` and then restore its attributes the way pickle
+    does: through ``__setstate__`` when the instance has one, and otherwise
+    by updating its ``__dict__``. The default ``__getstate__`` of an object
+    whose class defines ``__slots__`` returns a ``(dict_state, slots_state)``
+    tuple instead of a dict, and the slot values are set one by one.
+    """
+    if attrs is None:
+        return
+    state = attrs.construct()
+    if state is None:
+        return
+    if hasattr(instance, "__setstate__"):
+        instance.__setstate__(state)
+        return
+    slots_state = None
+    if isinstance(state, tuple) and len(state) == 2:
+        state, slots_state = state
+    if state:
+        instance.__dict__.update(state)
+    if slots_state:
+        for name, value in slots_state.items():
+            setattr(instance, name, value)
+
+
 def dict_get_state(obj: Any, save_context: SaveContext) -> dict[str, Any]:
     res = {
         "__class__": obj.__class__.__name__,
@@ -58,6 +125,9 @@ def dict_get_state(obj: Any, save_context: SaveContext) -> dict[str, Any]:
         content[key] = get_state(value, save_context)
     res["content"] = content
     res["key_types"] = key_types
+    attrs = _get_attrs_state(obj, save_context)
+    if attrs is not None:
+        res["attrs"] = attrs
     return res
 
 
@@ -78,16 +148,25 @@ def __init__(
             for key, value in state["content"].items()
         }
         self.children = {"key_types": self.key_types, "content": self.content}
+        # the instance attributes of a subclass, see ``_get_attrs_state``
+        self.attrs = _get_attrs_tree(state, load_context, trusted)
+        if self.attrs is not None:
+            self.children["attrs"] = self.attrs
 
     def _construct(self):
-        content = gettype(self.module_name, self.class_name)()
-        # Make the dict available to children which refer back to it, see
+        cls: type[object] = gettype(self.module_name, self.class_name)
+        # As pickle does, the instance is created with ``__new__`` instead of
+        # through its constructor, filled in place and then given its
+        # attributes back. It is stored before its items are constructed, so
+        # that items which refer back to it get the instance, see
         # ``Node.construct``.
-        self._constructed = content
+        instance: Any = cls.__new__(cls)
+        self._constructed = instance
         key_types = self.key_types.construct()
         for k_type, (key, val) in zip(key_types, self.content.items()):
-            content[k_type(key)] = val.construct()
-        return content
+            instance[k_type(key)] = val.construct()
+        _set_attrs(instance, self.attrs)
+        return instance
 
 
 def defaultdict_get_state(obj: Any, save_context: SaveContext) -> dict[str, Any]:
@@ -102,6 +181,9 @@ def defaultdict_get_state(obj: Any, save_context: SaveContext) -> dict[str, Any]
     content["main"] = get_state(dict(obj), save_context)
     content["default_factory"] = get_state(obj.default_factory, save_context)
     res["content"] = content
+    attrs = _get_attrs_state(obj, save_context)
+    if attrs is not None:
+        res["attrs"] = attrs
     return res
 
 
@@ -113,7 +195,7 @@ def __init__(
         trusted: TrustedTypes | None = None,
     ) -> None:
         super().__init__(state, load_context, trusted)
-        self.trusted = ["collections.defaultdict"]
+        self.trusted = self._get_trusted(trusted, ["collections.defaultdict"])
         self.main = get_tree(
             state["content"]["main"],
             load_context,
@@ -124,10 +206,24 @@ def __init__(
             state["content"]["default_factory"], load_context, trusted=trusted
         )
         self.children = {"main": self.main, "default_factory": self.default_factory}
+        # the instance attributes of a subclass, see ``_get_attrs_state``
+        self.attrs = _get_attrs_tree(state, load_context, trusted)
+        if self.attrs is not None:
+            self.children["attrs"] = self.attrs
 
     def _construct(self):
-        instance = defaultdict(**self.main.construct())
+        cls: type[object] = gettype(self.module_name, self.class_name)
+        # Like ``DictNode``: the instance is created with ``__new__``, stored,
+        # given its default factory, filled in place and then given its
+        # attributes back. Pickle builds a defaultdict through its constructor
+        # with the default factory instead, which a subclass whose constructor
+        # takes other arguments does not accept.
+        instance: Any = cls.__new__(cls)
+        self._constructed = instance
         instance.default_factory = self.default_factory.construct()
+        for key, value in self.main.construct().items():
+            instance[key] = value
+        _set_attrs(instance, self.attrs)
         return instance
 
 
@@ -140,6 +236,9 @@ def list_get_state(obj: Any, save_context: SaveContext) -> dict[str, Any]:
     content = [get_state(value, save_context) for value in obj]
 
     res["content"] = content
+    attrs = _get_attrs_state(obj, save_context)
+    if attrs is not None:
+        res["attrs"] = attrs
     return res
 
 
@@ -156,22 +255,23 @@ def __init__(
             get_tree(value, load_context, trusted=trusted) for value in state["content"]
         ]
         self.children = {"content": self.content}
+        # the instance attributes of a subclass, see ``_get_attrs_state``
+        self.attrs = _get_attrs_tree(state, load_context, trusted)
+        if self.attrs is not None:
+            self.children["attrs"] = self.attrs
 
     def _construct(self):
-        content_type = gettype(self.module_name, self.class_name)
-        if content_type is not list:
-            # Subclasses are built from their items through their own
-            # constructor, as before, so there is no instance to hand out
-            # before it is complete: a reference back to a list subclass is not
-            # supported, and ``get_state`` refuses it when saving.
-            return content_type([item.construct() for item in self.content])
-
-        # Fill a plain list in place and make it available to children which
-        # refer back to it, see ``Node.construct``.
-        content: list[Any] = []
-        self._constructed = content
-        content.extend(item.construct() for item in self.content)
-        return content
+        cls: type[object] = gettype(self.module_name, self.class_name)
+        # As pickle does, the instance is created with ``__new__`` instead of
+        # through its constructor, filled in place and then given its
+        # attributes back. It is stored before its items are constructed, so
+        # that items which refer back to it get the instance, see
+        # ``Node.construct``.
+        instance: Any = cls.__new__(cls)
+        self._constructed = instance
+        instance.extend([item.construct() for item in self.content])
+        _set_attrs(instance, self.attrs)
+        return instance
 
 
 def set_get_state(obj: Any, save_context: SaveContext) -> dict[str, Any]:
@@ -182,6 +282,9 @@ def set_get_state(obj: Any, save_context: SaveContext) -> dict[str, Any]:
     }
     content = [get_state(value, save_context) for value in obj]
     res["content"] = content
+    attrs = _get_attrs_state(obj, save_context)
+    if attrs is not None:
+        res["attrs"] = attrs
     return res
 
 
@@ -198,22 +301,24 @@ def __init__(
             get_tree(value, load_context, trusted=trusted) for value in state["content"]
         ]
         self.children = {"content": self.content}
+        # the instance attributes of a subclass, see ``_get_attrs_state``
+        self.attrs = _get_attrs_tree(state, load_context, trusted)
+        if self.attrs is not None:
+            self.children["attrs"] = self.attrs
 
     def _construct(self):
-        content_type = gettype(self.module_name, self.class_name)
-        if content_type is not set:
-            # Subclasses are built from their items through their own
-            # constructor, as before, so there is no instance to hand out
-            # before it is complete: a reference back to a set subclass is not
-            # supported, and ``get_state`` refuses it when saving.
-            return content_type([item.construct() for item in self.content])
-
-        # Fill a plain set in place and make it available to children which
-        # refer back to it, see ``Node.construct``.
-        content: set[Any] = set()
-        self._constructed = content
-        content.update(item.construct() for item in self.content)
-        return content
+        cls: type[object] = gettype(self.module_name, self.class_name)
+        # Like ``ListNode``: the instance is created with ``__new__``, stored,
+        # filled in place and then given its attributes back. Pickle builds a
+        # set subclass through its constructor instead, which would leave no
+        # instance to hand out to items which refer back to it. The items are
+        # added with the built-in ``set.update``, as the constructor would, so
+        # that an overridden ``update`` does not run on a bare instance.
+        instance: Any = cls.__new__(cls)
+        self._constructed = instance
+        set.update(instance, [item.construct() for item in self.content])
+        _set_attrs(instance, self.attrs)
+        return instance
 
 
 def tuple_get_state(obj: Any, save_context: SaveContext) -> dict[str, Any]:
@@ -578,19 +683,7 @@ def _construct(self):
         # Make the instance available to attributes which refer back to it, see
         # ``Node.construct``.
         self._constructed = instance
-
-        attrs_node = self.attrs
-        if attrs_node is None:
-            # nothing more to do
-            return instance
-
-        attrs = attrs_node.construct()
-        if attrs is not None:
-            if hasattr(instance, "__setstate__"):
-                instance.__setstate__(attrs)
-            else:
-                instance.__dict__.update(attrs)
-
+        _set_attrs(instance, self.attrs)
         return instance
 
 

skops/io/_utils.py

@@ -227,16 +227,19 @@ def _get_state(obj, save_context: SaveContext):
     raise TypeError(f"Getting the state of type {type(obj)} is not supported yet")
 
 
-def _supports_circular_reference(value: Any, state: dict[str, Any]) -> bool:
-    """Whether a reference to ``value`` from inside its own state can be loaded.
+def _supports_circular_reference(state: dict[str, Any]) -> bool:
+    """Whether a reference to the object of ``state`` from inside it can be loaded.
 
     This mirrors the ``Node`` classes whose ``_construct`` registers the
     instance before constructing its children, see ``Node.construct``.
-    ``ListNode`` and ``SetNode`` only do so for plain lists and sets.
     """
-    if type(value) in (list, set):
-        return True
-    return state["__loader__"] in ("DictNode", "ObjectNode")
+    return state["__loader__"] in (
+        "DictNode",
+        "DefaultDictNode",
+        "ListNode",
+        "SetNode",
+        "ObjectNode",
+    )
 
 
 def get_state(value, save_context: SaveContext) -> dict[str, Any]:
@@ -264,9 +267,7 @@ def get_state(value, save_context: SaveContext) -> dict[str, Any]:
     save_context.in_progress[__id__] = False
     try:
         res = _get_state(value, save_context)
-        if save_context.in_progress[__id__] and not _supports_circular_reference(
-            value, res
-        ):
+        if save_context.in_progress[__id__] and not _supports_circular_reference(res):
             raise UnsupportedTypeException(
                 f"Objects of type {type(value).__name__} which contain a"
                 " reference to themselves are not supported yet."