You need to enable JavaScript to run this app.
最新活动
大模型
产品
解决方案
定价
生态与合作
支持与服务
开发者
了解我们

如何在包含对象上调用__setstate__实现pickle协同序列化

解决自定义Pickle序列化中类B的状态恢复问题

这个报错的核心原因很明确:反序列化B的实例时,新建的空对象还没有self.a属性,你直接调用self.a.__setstate__()自然会触发AttributeError。Pickle反序列化的流程是先创建类的空实例(不执行__init__),再调用__setstate__传入保存的状态,所以我们需要在__setstate__中先完成A实例的初始化,再恢复它的状态。

修复方案

我们需要修改B类的__setstate__方法,先创建A的实例,再调用它的__setstate__恢复状态,同时设置B自身的属性。具体代码如下:

import pickle

class A:
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z
    def __eq__(self, other):
        return self.x == other.x and self.y == other.y and self.z == other.z
    def __getstate__(self):
        return self.x, self.y, self.z
    def __setstate__(self, state):
        self.x, self.y, self.z = state

class B:
    def __init__(self):
        self.a = A(1, 2, 3)
        self.b = 4
    def __getstate__(self):
        # 保存A的自定义状态和B自身的属性,这部分是正确的
        return self.a.__getstate__(), self.b
    def __setstate__(self, state):
        a_state, b_value = state
        # 1. 先创建A的实例(临时初始值不影响,之后会被__setstate__覆盖)
        self.a = A(0, 0, 0)
        # 2. 调用A的__setstate__恢复它的原始状态
        self.a.__setstate__(a_state)
        # 3. 恢复B自身的属性
        self.b = b_value
    def __eq__(self, other):
        return self.a == other.a and self.b == other.b

# 测试A的序列化/反序列化
a = A(1, 2, 3)
a_save = pickle.dumps(a)
assert a == pickle.loads(a_save)

# 测试B的序列化/反序列化
b = B()
b_save = pickle.dumps(b)
assert b == pickle.loads(b_save)
print("所有测试通过!")

优化说明

如果A类的__init__参数正好匹配__getstate__返回的状态元组,你还可以简化__setstate__,直接通过状态创建A实例,省去调用__setstate__的步骤:

def __setstate__(self, state):
    a_state, self.b = state
    # 直接用A的状态元组初始化实例,效果和调用__setstate__一致
    self.a = A(*a_state)

这种方式更高效,但前提是A的__init____setstate__的逻辑完全一致。如果A的__setstate__有额外的自定义逻辑(比如初始化动态生成的属性),还是需要保留调用__setstate__的写法。

内容的提问来源于stack exchange,提问作者Gil Hamilton

火山引擎 最新活动