使用更改的模块路径对 python 对象进行 unpickling

Question

我正在尝试将同事构建的项目Project A 集成到另一个 python 项目中。现在这位同事在他的代码中没有使用相对导入，而是使用了

from packageA.moduleA import ClassA
from packageA.moduleA import ClassB

因此使用cPickle 腌制课程。为了整洁，我想隐藏他 (Project A) 在我的项目中构建的包。然而，这改变了packageA 中定义的类的路径。没问题，我将使用重新定义导入

from ..packageA.moduleA import ClassA
from ..packageA.moduleA import ClassB

但现在取消腌制类失败并显示以下消息

    with open(fname) as infile: self.clzA = cPickle.load(infile)
ImportError: No module named packageA.moduleA

那么为什么cPickle 显然看不到模块定义。我需要将packageA 的根添加到系统路径吗？这是解决问题的正确方法吗？

cPickled 文件看起来像

ccopy_reg
_reconstructor
p1
(cpackageA.moduleA
ClassA
p2
c__builtin__
object
p3
NtRp4

旧的项目层次结构是这样的

packageA/
    __init__.py
    moduleA.py
    moduleB.py
packageB/
    __init__.py
    moduleC.py
    moduleD.py

我想把所有这些都放入WrapperPackage

MyPackage/
.. __init__.py
.. myModuleX.py
.. myModuleY.py
WrapperPackage/
.. __init__.py
.. packageA/
   .. __init__.py
   .. moduleA.py
   .. moduleB.py
.. packageB/
   .. __init__.py
   .. moduleC.py
   .. moduleD.py

Answer 1

您需要为 pickle 导入创建一个别名才能工作； WrapperPackage包的__init__.py文件如下：

from .packageA import * # Ensures that all the modules have been loaded in their new locations *first*.
from . import packageA  # imports WrapperPackage/packageA
import sys
sys.modules['packageA'] = packageA  # creates a packageA entry in sys.modules

您可能需要创建其他条目：

sys.modules['packageA.moduleA'] = moduleA
# etc.

现在 cPickle 将在它们的旧位置再次找到 packageA.moduleA 和 packageA.moduleB。

之后您可能需要重新编写 pickle 文件，届时将使用新的模块位置。上面创建的附加别名应确保有问题的模块具有新的位置名称，以便在再次编写类时使用 cPickle。

Answer 2

除了@MartinPieters 回答之外，另一种方法是定义cPickle.Unpickler 类的find_global 方法，或者扩展pickle.Unpickler 类。

def map_path(mod_name, kls_name):
    if mod_name.startswith('packageA'): # catch all old module names
        mod = __import__('WrapperPackage.%s'%mod_name, fromlist=[mod_name])
        return getattr(mod, kls_name)
    else:
        mod = __import__(mod_name)
        return getattr(mod, kls_name)

import cPickle as pickle
with open('dump.pickle','r') as fh:
    unpickler = pickle.Unpickler(fh)
    unpickler.find_global = map_path
    obj = unpickler.load() # object will now contain the new class path reference

with open('dump-new.pickle','w') as fh:
    pickle.dump(obj, fh) # ClassA will now have a new path in 'dump-new'

pickle 和cPickle 的过程的更详细说明可以找到here。

Answer 3

一种可能的解决方案是直接编辑 pickle 文件（如果您有权访问）。我遇到了更改模块路径的同样问题，并且我将文件保存为 pickle.HIGHEST_PROTOCOL 所以理论上它应该是二进制的，但是模块路径以纯文本形式位于 pickle 文件的顶部。因此，我只是对旧模块路径的所有实例进行了查找替换，并使用新的，瞧，它们正确加载了。

我确信这个解决方案并不适合所有人，特别是如果你有一个非常复杂的腌制对象，但它是一种快速而肮脏的数据修复方法，对我有用！

Answer 4

这是我灵活解酸的基本模式——通过明确且快速的转换图——因为除了与酸洗相关的原始数据类型之外，通常只有几个已知的类。这还可以保护 unpickling 免受错误或恶意构造的数据的影响，毕竟这些数据可以在简单的pickle.load() 上执行任意 python 代码（！）（无论是否有容易出错的 sys.modules 摆弄）。

Python 2 和 3：

from __future__ import print_function
try:    
    import cPickle as pickle, copy_reg as copyreg
except: 
    import pickle, copyreg

class OldZ:
    a = 1
class Z(object):
    a = 2
class Dangerous:
    pass   

_unpickle_map_safe = {
    # all possible and allowed (!) classes & upgrade paths    
    (__name__, 'Z')         : Z,    
    (__name__, 'OldZ')      : Z,
    ('old.package', 'OldZ') : Z,
    ('__main__', 'Z')       : Z,
    ('__main__', 'OldZ')    : Z,
    # basically required
    ('copy_reg', '_reconstructor') : copyreg._reconstructor,    
    ('__builtin__', 'object')      : copyreg._reconstructor,    
    }

def unpickle_find_class(modname, clsname):
    print("DEBUG unpickling: %(modname)s . %(clsname)s" % locals())
    try: 
        return _unpickle_map_safe[(modname, clsname)]
    except KeyError:
        raise pickle.UnpicklingError(
            "%(modname)s . %(clsname)s not allowed" % locals())
if pickle.__name__ == 'cPickle':  # PY2
    def SafeUnpickler(f):
        u = pickle.Unpickler(f)
        u.find_global = unpickle_find_class
        return u
else:  # PY3 & Python2-pickle.py
    class SafeUnpickler(pickle.Unpickler):  
        find_class = staticmethod(unpickle_find_class)

def test(fn='./z.pkl'):
    z = OldZ()
    z.b = 'teststring' + sys.version
    pickle.dump(z, open(fn, 'wb'), 2)
    pickle.dump(Dangerous(), open(fn + 'D', 'wb'), 2)
    # load again
    o = SafeUnpickler(open(fn, 'rb')).load()
    print(pickle, "loaded:", o, o.a, o.b)
    assert o.__class__ is Z
    try: 
        raise SafeUnpickler(open(fn + 'D', 'rb')).load() and AssertionError
    except pickle.UnpicklingError: 
        print('OK: Dangerous not allowed')

if __name__ == '__main__':
    test()