如何用dill存储Python子模块输入以实现跨环境调试？

我开发了一个Python模块MyModule，用户会将其作为git submodule引入，导入方式如下：

from MyModule.MyModuleSubfolder import ClassA
from MyModule.MyModuleSubfolder import ClassB
from MyModule.MyModuleOtherSubfolder import useful_function

用户传给useful_function的是包含字典、列表、大型numpy数组的复杂字典。我希望在useful_function里用dill存储这个输入，方便调试读取。

我实现的核心代码如下（省略文件IO逻辑）：

...
input_dict['file_path'] = "folder\input_file.pkl"
...

def useful_function(input_dict, load_flag=False):
    if load_flag:
        input_dict = dill.load(input_dict['file_path'], ...)
    else:
        dill.dump(input_dict, ...)

这个实现不管是我把MyModule作为独立项目运行，还是用户作为git submodule使用时，读写都正常。但当我在独立项目中加载用户作为子模块运行时生成的input_file.pkl时，会报错：ModuleNotFoundError: No module named 'MyModule'。

类似问题的“mainify”方案不适用：一是我只需要存储复杂字典，不需要序列化MyModule内的类，但pickle/dill会保存主环境的所有导入；二是我无法控制用户的导入行为，没法预先处理所有可能的导入。求更优解决方案？

方法1：序列化前剥离无关模块引用 #

既然不需要序列化MyModule内的类，可在dump前递归清理input_dict，只保留基础类型、numpy数组等你需要的对象，避免dill捕获无关的模块引用：

import dill
import numpy as np

def clean_for_serialization(obj):
    if isinstance(obj, dict):
        return {k: clean_for_serialization(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [clean_for_serialization(item) for item in obj]
    elif isinstance(obj, np.ndarray):
        return obj
    # 按需添加其他需要保留的基础类型
    elif isinstance(obj, (tuple, str, int, float, bool)):
        return obj
    else:
        # 若输入中不应存在自定义类，可抛出警告或直接过滤
        raise ValueError(f"Unexpected object type {type(obj)} found in input_dict")

def useful_function(input_dict, load_flag=False):
    if load_flag:
        with open(input_dict['file_path'], 'rb') as f:
            input_dict = dill.load(f)
    else:
        cleaned_dict = clean_for_serialization(input_dict)
        with open(input_dict['file_path'], 'wb') as f:
            dill.dump(cleaned_dict, f)
    return input_dict

方法2：加载时自定义模块映射 #

利用dill的模块映射机制，在加载pickle文件时，将用户环境中找不到的MyModule路径映射到当前环境的对应模块，无需修改dump逻辑：

import dill
import sys
from MyModule import MyModuleSubfolder, MyModuleOtherSubfolder

def useful_function(input_dict, load_flag=False):
    if load_flag:
        # 映射具体类（按需添加）
        dill._dill._reverse_typemap.update({
            'MyModule.MyModuleSubfolder.ClassA': MyModuleSubfolder.ClassA,
            'MyModule.MyModuleSubfolder.ClassB': MyModuleSubfolder.ClassB,
        })
        # 通用映射：将MyModule指向当前环境的模块
        sys.modules['MyModule'] = sys.modules[__name__.split('.')[0]]
        
        with open(input_dict['file_path'], 'rb') as f:
            input_dict = dill.load(f)
    else:
        with open(input_dict['file_path'], 'wb') as f:
            dill.dump(input_dict, f)
    return input_dict

方法3：替换为无模块依赖的序列化方式 #

如果输入仅包含基础类型和numpy数组，可使用numpy自带的np.savez或JSON（配合numpy数组转换），彻底避免pickle/dill的模块依赖问题：

import numpy as np

def useful_function(input_dict, load_flag=False):
    if load_flag:
        with np.load(input_dict['file_path'], allow_pickle=True) as data:
            input_dict = data['input_dict'].item()
    else:
        np.savez(input_dict['file_path'], input_dict=input_dict)
    return input_dict

若选择JSON，需额外处理numpy数组与列表的转换，适合需要跨语言读取的场景。

内容的提问来源于stack exchange，提问作者Shim'on