如何在copyreg中注册dill以序列化不可pickle的模块？

场景复现 #

我遇到的场景是正在使用的Python包需要序列化我的代码，简化示例如下：

constraint.py #

manifest = []

constraint_loader.py #

import importlib

class constraint_loader:
    def __init__(self):
        self.constraints = importlib.import_module('constraint')

test.py #

import pickle
import constraint_loader

pickle.dumps(constraint_loader.constraint_loader())

运行上述代码会报错：

% python test.py
Traceback (most recent call last):
  File "test.py", line 4, in <module>
    pickle.dumps(constraint_loader.constraint_loader())
TypeError: cannot pickle 'module' object

改用dill可以正常序列化：
修改后的test.py：

import dill
import constraint_loader

dill.dumps(constraint_loader.constraint_loader())

尝试的错误方案 #

我尝试通过copyreg注册dill来让pickle支持模块序列化，修改后的constraint_loader.py如下：

import importlib
import copyreg
import dill
from types import ModuleType

class constraint_loader:
    def __init__(self):
        self.constraints = importlib.import_module('constraint')

copyreg.pickle(ModuleType, dill.dumps)

运行后依然报错：

% python test.py
Traceback (most recent call last):
  File "test.py", line 4, in <module>
    pickle.dumps(constraint_loader.constraint_loader())
_pickle.PicklingError: __reduce__ must return a string or tuple

正确解决方法 #

copyreg.pickle的第二个参数需要是一个返回可调用对象和参数元组的函数，而不是直接传入dill.dumps。我们可以包装一个适配函数，利用dill的序列化逻辑来生成pickle能识别的还原信息：

修改constraint_loader.py如下：

import importlib
import copyreg
import dill
from types import ModuleType

def pickle_module(module):
    # 用dill序列化模块，再封装成可以被pickle还原的逻辑
    return dill.loads, (dill.dumps(module),)

class constraint_loader:
    def __init__(self):
        self.constraints = importlib.import_module('constraint')

copyreg.pickle(ModuleType, pickle_module)

现在再运行原来的test.py，pickle就能正常序列化包含模块对象的constraint_loader实例了。

原理说明 #

• copyreg.pickle要求注册的函数返回一个二元组：第一个元素是还原对象的可调用函数，第二个元素是传给该函数的参数元组。
• 这里我们用dill.dumps把模块序列化成字节串，然后用dill.loads作为还原函数，把字节串作为参数传入，这样就能在反序列化时还原出原模块对象。

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