动态加载临时模块类的实例pickle序列化失败问题及优化方案咨询

报错根本原因 #

pickle序列化自定义类实例时，不会存储完整的类定义，仅会存储「类所属模块名 + 类名 + 实例属性数据」。序列化阶段就会校验对应模块是否可导入，你的临时模块对应的.py文件在with块执行结束后就被自动删除，且动态加载的模块默认没有注册到sys.modules全局索引中，pickle找不到对应模块就会抛出该错误。

方案1：注册临时模块到sys.modules（最简便，完全不用改save方法） #

加载完临时模块后，手动将其加入sys.modules全局索引，只要当前进程不退出，pickle就能正常找到该模块完成序列化：

import sys
import tempfile
import importlib.util

with tempfile.NamedTemporaryFile(suffix='.py') as temp:
    temp.write(code.encode('utf-8'))
    spec = importlib.util.spec_from_file_location('a_temp_module', temp.name)
    temp_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(temp_module)
    # 新增这一行完成模块注册
    sys.modules['a_temp_module'] = temp_module
    special_snowflake_object = temp_module.SpecialSnowflake()

# 之后直接调用save即可，不需要再重载模块
special_snowflake_object.save(path_to_pickle)

额外注意：你现有save/load方法里的open模式写错了，pickle是二进制协议，需要用wb/rb而非w/r，否则会出现编码报错，修改后：

def save(self, file_path: str):
    with open(file_path, 'wb') as fp:
        pickle.dump(self, fp)

@staticmethod
def load(file_path: str):
    with open(file_path, 'rb') as fp:
        return pickle.load(fp)

方案2：自定义__reduce__方法，完全脱离模块依赖 #

如果需要跨进程序列化也能正常工作，可以给SpecialSnowflake类实现__reduce__方法，自定义序列化/反序列化逻辑，不需要依赖临时模块在当前进程的注册状态：
首先在全局作用域定义重构函数：

def _reconstruct_snowflake(code: str, instance_attrs: dict):
    import tempfile
    import importlib.util
    with tempfile.NamedTemporaryFile(suffix='.py') as temp:
        temp.write(code.encode('utf-8'))
        spec = importlib.util.spec_from_file_location('a_temp_module', temp.name)
        temp_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(temp_module)
        obj = temp_module.SpecialSnowflake.__new__(temp_module.SpecialSnowflake)
        obj.__dict__.update(instance_attrs)
        return obj

然后在SpecialSnowflake类中添加__reduce__方法，注意提前把类的源码code存为实例属性：

class SpecialSnowflake:
    def __init__(self, code):
        self.code = code
        # 其他初始化逻辑
    
    def __reduce__(self):
        # 返回重构函数 + 所需参数
        return (_reconstruct_snowflake, (self.code, self.__dict__))
    
    # 其余原有方法无需修改

这个方案下save方法可以随时调用，甚至跨进程传输pickle文件也能正常反序列化。

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