腌制具有 __slots__ 的冻结数据类

Posted 2023-02-23

【中文标题】腌制具有 __slots__ 的冻结数据类

【英文标题】：Pickle a frozen dataclass that has __slots__

【发布时间】：2019-08-13 21:03:15

【问题描述】：

如何使用__slots__ 腌制冻结数据类的实例？例如，以下代码在 Python 3.7.0 中引发异常：

import pickle
from dataclasses import dataclass

@dataclass(frozen=True)
class A:
  __slots__ = ('a',)
  a: int

b = pickle.dumps(A(5))
pickle.loads(b)

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<string>", line 3, in __setattr__
dataclasses.FrozenInstanceError: cannot assign to field 'a'

如果我删除 frozen 或 __slots__，这将有效。这只是一个错误吗？

【参考方案1】：

问题来自pickle在设置插槽状态时使用实例的__setattr__方法。

默认__setstate__在load_build中定义_pickle.c line 6220。

对于状态dict中的项目，直接更新实例__dict__：

 if (PyObject_SetItem(dict, d_key, d_value) < 0)

而对于 slotstate dict 中的项目，则使用实例的 __setattr__：

if (PyObject_SetAttr(inst, d_key, d_value) < 0)

现在因为实例被冻结，__setattr__ 在加载时引发FrozenInstanceError。

要避免这种情况，您可以定义自己的__setstate__ 方法，该方法将使用object.__setattr__，而不是实例的__setattr__。

docs 对此给出了某种警告：

使用frozen=True 时会有微小的性能损失：__init__() 不能使用简单赋值来初始化字段，而必须使用 object.__setattr__()。

在您的情况下，将__getstate__ 定义为实例__dict__ 始终是None 也可能很好。如果不这样做，__setstate__ 的 state 参数将是一个元组 (None, 'a': 5)，第一个值是实例的 __dict__ 的值，第二个是 slotstate dict。

import pickle
from dataclasses import dataclass

@dataclass(frozen=True)
class A:
    __slots__ = ('a',)
    a: int

    def __getstate__(self):
        return dict(
            (slot, getattr(self, slot))
            for slot in self.__slots__
            if hasattr(self, slot)
        )

    def __setstate__(self, state):
        for slot, value in state.items():
            object.__setattr__(self, slot, value) # <- use object.__setattr__


b = pickle.dumps(A(5))
pickle.loads(b)

我个人不会将其称为错误，因为酸洗过程被设计为灵活，但仍有增强功能的空间。酸洗协议的修订可以在未来解决这个问题。除非我遗漏了一些东西并且除了微小的性能损失，对所有插槽使用PyObject_GenericSetattr 可能是一个合理的解决方法？

【讨论】：

我更新了之前删除的答案。这一切都归结为创建自己的__setstate__ 方法以避免调用实例的__setattr__。

你知道为什么如果类没有插槽它会起作用吗？在那种情况下它不使用 setattr 来初始化吗？

@Arne 好点，我刚刚看了pickle的源代码，似乎其他属性是通过直接修改__dict__来处理的，而插槽使用__setattr__。我会深入了解并更新。

我有reported 这个问题并向 Python 错误跟踪器推荐了您的解决方法。希望它能成为标准库。

@drhagen 太好了，有机会我会尝试使用PyObject_GenericSetattr 进行编译。手指交叉！

【参考方案2】：

从 Python 3.10.0 开始，这有效，但前提是您在数据类装饰器中通过 slots=True 指定插槽。手动指定 __slots__ 时，它不起作用，并且可能永远不会起作用。

import pickle
from dataclasses import dataclass

@dataclass(frozen=True, slots=True)
class A:
  a: int

b = pickle.dumps(A(5))
pickle.loads(b)  # A(a=5)

【参考方案3】：

如果您只需要类是可散列的，您可以使用unsafe_hash=True 选项强制生成__hash__ 函数。你不会得到不变性保证，但无论如何，python 中的不变性是不可能的。

Relevant python documentation 状态：

虽然不推荐，但您可以强制dataclass() 使用unsafe_hash=True 创建__hash__() 方法。如果您的类在逻辑上是不可变的，但仍然可以发生变异，则可能会出现这种情况。这是一个专门的用例，应该仔细考虑。

import pickle
from dataclasses import dataclass

@dataclass(unsafe_hash=True)
class A:
    __slots__ = ('a',)
    a: int

b = pickle.dumps(A(5))
hash(pickle.loads(b))  # works and can hash!