pickle跨模块进行序列化的解决方案

问题背景

在 Python 中，pickle 是一个常用的序列化工具，可以把对象转为字节流，再在另一端还原。但是它有一个关键限制：反序列化时必须能在相同的 module 和 class name 下找到对应的类。通常情况下，我们会在同一个代码环境中进行pickle的dumps和load，但是如果涉及到一些非常规的任务，就可能遇到这个序列化和反序列化的问题。

我们首先描述一下目前的问题。考虑我们有两个环境，分别是发送端和接收端，在发送端中，存在一个a.b.c.A的类，而在接收端中存在另一个x.y.z.B的类，这两个类的定义是完全相同的，区别在于module和class name不同。而现在我们要做的就是在发送端中使用pickle dumps出一个A的字节流，然后在接收端从字节流中恢复成B。用代码表述如下：

# 发送端的定义
# module: a/b/c.py
class A:
    def __init__(self, data: int):
        self.data = data
a = A(123)
data = pickle.dumps(a)
with open('data.pkl', "wb") as f:
    f.write(data)

# 接收端的定义
# module: x/y/z.py
class B:
    def __init__(self, data: int):
        self.data = data
with open('data.pkl', 'rb') as f:
    obj = pickle.load(f)
assert isinstance(obj, B)

如果直接进行跨模块的序列化和反序列化，会在接收端报错：

ModuleNotFoundError: No module named 'a'

接下来我们分析一下为什么会失败。pickle 在序列化对象时，会记录对象的__module__ 和 __name__，这些信息是存放在字节流中的。而在反序列化的时候，就需要将将它生成回对应的module和name。然而在接收端环境中并不存在对应的a.b.c.A，所以会报错找不到Module。

解决方案

我们可以在不同的层次上解决这个问题，对应不同的修改方式。核心原理都是需要在字节流中记录可以被接收端正确import的module和name。

第一种方式是使用一个通用的中间结构来存储信息，例如dict。我们在发送端将A的内容转化成dict之后再进行序列化，而在接收端则通过dict来初始化对应的B。

# 发送端代码修改
from a.b.c import A
import pickle


a = A(123)
data = pickle.dumps(a.__dict__)  # 利用中间结构dict来存储信息
with open("data.pkl", "wb") as f:
    f.write(data)

# 接收端代码修改
from x.y.z import B
import pickle

with open('data.pkl', 'rb') as f:
    obj = pickle.load(f)
    obj = B(**obj)  # 利用dict来初始化对应的类

assert isinstance(obj, B)

这种方式需要同时修改发送端和接收端的代码。

第二种方式则是直在发送端构造出新的x.y.z.B，module和name与接收端保持一致即可。

这种方式对发送端代码整体架构影响较大。

第三种方式则无需对接收端代码进行修改，并且不需要对发送端代码整体架构进行调整。因为pickle反序列化需要对应的module和name，我们可以在发送端进行伪造：

A.__module__ = 'x.y.z'
A.__qualname__ = 'B'

• 这里需要注意__name__和__qualname__的区别：name表示对象的短名称，对类来说就是类名，对函数来说就是函数名；而qualname表示完整的作用域路径，包含了嵌套关系。
• 在pickle中，实际上依赖的是__module__和__qualname__

当然仅仅这样是不行的，因为在dumps序列化的时候，也需要对应的信息，它会报错如下：

_pickle.PicklingError: Can't pickle <class 'x.y.z.B'>: import of module 'x.y.z' failed

这个错误是说在发送端找不到x.y.z，确实在发送端并不存在这样的module。因此接下来的一步就是动态的构造出这个module，并且无需调整代码的组织架构。我们可以通过types.ModuleType来伪造模块，并将其注入到sys.modules里面，示例如下：

import sys
import types

sys.modules["x"] = types.ModuleType("x")
sys.modules["x.y"] = types.ModuleType("x.y")
sys.modules["x.y.z"] = types.ModuleType("x.y.z")

from a.b import c

fake_xyz_module = sys.modules["x.y.z"]
fake_xyz_module.__dict__.update(c.__dict__)  # 将module中的信息同步更新到伪造module中
fake_xyz_module.B = c.A  # 同步伪造一个fake B

完整代码如下：

# 仅修改发送端代码
from a.b.c import A
import pickle

# 动态构造module
import sys
import types

sys.modules["x"] = types.ModuleType("x")
sys.modules["x.y"] = types.ModuleType("x.y")
sys.modules["x.y.z"] = types.ModuleType("x.y.z")

# 更新伪造的module中的类信息
from a.b import c

fake_xyz_module = sys.modules["x.y.z"]
fake_xyz_module.__dict__.update(c.__dict__)
fake_xyz_module.B = c.A

A.__module__ = "x.y.z"
A.__qualname__ = 'B'

# pickle序列化
a = A(123)
data = pickle.dumps(a)
with open("data.pkl", "wb") as f:
    f.write(data)

这样我们就只对发送端代码进行修改，并且无需修改发送端代码的整体架构。当然无论如何，这种方式只能在类结构兼容时使用。如果类定义差异很大，则很难通过pickle直接进行跨模块的序列化和反序列化。

---

「NOTE」最后，这里还提供一个简单的递归动态生成module的工具函数：

import types
import sys

def create_module_hierarchy(path: str):
    """
    create a module hierarchy in sys.modules
    """
    parts = path.split(".")
    for i in range(1, len(parts) + 1):
        sub_path = ".".join(parts[:i])
        if sub_path not in sys.modules:
            mod = types.ModuleType(sub_path)
            sys.modules[sub_path] = mod
            if i > 1:
                parent_path = ".".join(parts[:i - 1])
                setattr(sys.modules[parent_path], parts[i - 1], mod)

create_module_hierarchy("x.y.z")