从SQLite加载Pickle对象时触发AttributeError错误求助

场景复现 #

将自定义类F的对象通过Pickle序列化后存入SQLite数据库，存储代码（pickle.py）：

import sqlite3
import pickle
import pandas as pd

sqlite3.register_converter("pickle", pickle.loads)
sqlite3.register_adapter(list, pickle.dumps)
sqlite3.register_adapter(set, pickle.dumps)

class F:
    a = None
    b = None
    def __init__(self) -> None:
        pass

df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})
f = F()
f.a = df
f.b = df.columns

data = pickle.dumps(f, protocol=pickle.HIGHEST_PROTOCOL)

sqliteConnection = sqlite3.connect('SQLite_Python.db')
cursor = sqliteConnection.cursor()
print("Successfully Connected to SQLite")
DATA = sqlite3.Binary(data)
sqlite_insert_query = """INSERT INTO PICKLES1 (INTEGRATION_NAME, DATA) VALUES ('James',?)"""

resp = cursor.execute(sqlite_insert_query,(DATA,))
sqliteConnection.commit()

从数据库读取并反序列化代码（retrieve_pickle.py）：

import sqlite3
import pickle

conn = sqlite3.connect('SQLite_Python.db')
cur = conn.cursor()
cur.execute("SELECT DATA FROM PICKLES1 where INTEGRATION_NAME='James'")
df = None
rows = cur.fetchall()
for r in rows[0]:
    print(type(r)) # 输出 <class 'bytes'>
    df = pickle.loads(r)

执行时触发错误：

File "/Users/ETC/Work/pickle_work/picklertry.py", line 34, in select_all_tasks
    df = pickle.loads(r)
AttributeError: Can't get attribute 'F' on <module '__main__' from '/Users/rusab1/Work/pickle_work/picklertry.py'>

错误原因 #

Pickle序列化自定义对象时，仅记录类的模块名称和类名，不会保存类的完整定义。序列化时，类F属于pickle.py的__main__模块；反序列化时，当前运行的retrieve_pickle.py的__main__模块中没有定义F类，因此无法找到对应类完成实例化。

解决方法 #

方法1：将自定义类放到独立模块（推荐） #

把F类单独存入可导入的模块（比如my_classes.py）：

# my_classes.py
class F:
    a = None
    b = None
    def __init__(self) -> None:
        pass

在序列化和反序列化代码中统一导入该类：

• pickle.py中添加：

  from my_classes import F
  

• retrieve_pickle.py中添加：

  from my_classes import F
  

这样序列化与反序列化时，F类都来自同一模块my_classes，Pickle可正确匹配类定义。

方法2：在反序列化模块中手动注册类 #

若无法拆分模块，可在retrieve_pickle.py中定义与pickle.py完全一致的F类，或导入后绑定到当前__main__模块：

# retrieve_pickle.py
import sys
# 定义与pickle.py中完全相同的F类
class F:
    a = None
    b = None
    def __init__(self) -> None:
        pass

# 若已从其他模块导入，也可直接绑定到__main__
# from some_module import F
# sys.modules['__main__'].F = F

方法3：自定义Pickle Unpickler #

重写Unpickler的find_class方法，指定F类的正确来源：

# retrieve_pickle.py
import pickle
import io
from my_classes import F

class CustomUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == '__main__' and name == 'F':
            return F
        return super().find_class(module, name)

# 使用自定义Unpickler加载数据
for r in rows[0]:
    df = CustomUnpickler(io.BytesIO(r)).load()

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