为numpy.ndarray子类添加可pickle序列化的属性

你碰到的这个问题是numpy ndarray子类的典型pickle序列化坑——numpy默认的序列化逻辑只会处理数组的核心数据（形状、dtype、缓冲区这些），不会自动保存你添加的自定义属性（比如这里的_csys），所以反序列化后这个属性就丢失了，导致访问csys时触发AttributeError。

问题根源 #

当你pickle一个ndarray子类实例时，numpy内部的序列化机制并不会主动收集你添加的实例属性。反序列化重建对象时，__array_finalize__虽然尝试从原对象获取_csys，但pickle重建过程中并没有传递这个属性，所以新实例的_csys根本没被初始化。

解决方案：实现__getstate__和__setstate__ #

要让pickle保存你的自定义属性，只需要在子类中添加__getstate__方法来定义要保存的状态（包含数组本身和自定义属性），再用__setstate__来恢复这些状态。

修改后的完整代码：

import numpy as np

class Point(np.ndarray):
    def __new__(cls, arr, csys=None):
        obj = np.asarray(arr, dtype=np.float64).view(cls)
        obj._csys = csys
        return obj
    
    def __array_finalize__(self, obj):
        if obj is None:
            return
        self._csys = getattr(obj, '_csys', None)
    
    @property
    def csys(self):
        print('Getting .csys')
        return self._csys
    
    @csys.setter
    def csys(self, csys):
        print('Setting .csys')
        self._csys = csys
    
    # 新增：告诉pickle要保存的状态
    def __getstate__(self):
        # 复制实例的属性字典，再加上numpy数组的核心状态
        state = self.__dict__.copy()
        state['numpy_state'] = super().__reduce__()[2]
        return state
    
    # 新增：反序列化时恢复状态
    def __setstate__(self, state):
        # 先恢复numpy数组的核心状态
        super().__setstate__(state.pop('numpy_state'))
        # 恢复自定义属性
        self.__dict__.update(state)

测试验证 #

用你原来的测试代码跑一遍：

pt = Point([1, 2, 3])
pt.csys = 'cmm'
print("pt.csys:", pt.csys)
# 序列化、反序列化后检查
import pickle
pklstr = pickle.dumps(pt)
ppt = pickle.loads(pklstr)
print("ppt.csys:", ppt.csys)

输出结果会是：

Setting .csys
Getting .csys
pt.csys: cmm
Getting .csys
ppt.csys: cmm

完美解决属性丢失的问题。

另一种可选方案：重写__reduce__ #

如果你更熟悉pickle的底层逻辑，也可以重写__reduce__方法来包含自定义属性：

def __reduce__(self):
    # 获取父类的reduce返回值（类型、构造参数、状态）
    base_reduce = super().__reduce__()
    # 把自定义属性追加到状态元组中
    new_state = base_reduce[2] + (self._csys,)
    return (base_reduce[0], base_reduce[1], new_state)

def __setstate__(self, state):
    # 先取出自定义属性，剩下的交给父类处理
    self._csys = state[-1]
    super().__setstate__(state[:-1])

这种方式也能达到同样效果，但__getstate__/__setstate__的写法更直观，后期维护起来更简单。

内容的提问来源于stack exchange，提问作者subnivean