如何在使用Python的Pickle加载时将模块.Class()替换为本地定义的Class()?

Question

class Bar:
    pass

foo_module = imp.new_module('foo')
foo_module.Bar = Bar
sys.modules['foo'] = foo_module

import foo
print foo.Bar()

我有一个用 pickle 保存的数据，这里面有一堆 foo.Bar() 对象。我想把它们恢复出来，但 Bar() 这个类的定义在我尝试恢复的同一个文件里，而不在 foo 模块里。所以，pickle 就报错说找不到 foo 模块。

我试着用类似下面的方式来引入 foo 模块：

import imp, sys

这样是可以的，但当我在之后添加：

import pickle
p = pickle.load(open("my-pickle.pkl"))

我就收到了一个友好的错误提示：

Traceback (most recent call last):
  File "pyppd.py", line 69, in 
    ppds = loads(ppds_decompressed)
  File "/usr/lib/python2.6/pickle.py", line 1374, in loads
    return Unpickler(file).load()
  File "/usr/lib/python2.6/pickle.py", line 858, in load
    dispatch[key](self)
  File "/usr/lib/python2.6/pickle.py", line 1069, in load_inst
    klass = self.find_class(module, name)
  File "/usr/lib/python2.6/pickle.py", line 1124, in find_class
    __import__(module)
  File "/tmp/test.py", line 69, in 
    p = pickle.load(open("my-pickle.pkl"))
  File "/usr/lib/python2.6/pickle.py", line 1374, in loads
    return Unpickler(file).load()
  File "/usr/lib/python2.6/pickle.py", line 858, in load
    dispatch[key](self)
  File "/usr/lib/python2.6/pickle.py", line 1069, in load_inst
    klass = self.find_class(module, name)
  File "/usr/lib/python2.6/pickle.py", line 1124, in find_class
    __import__(module)
ImportError: No module named foo

有没有什么想法？

Answer 1

class Bar:
    pass

class MyUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == "foo" and name == "Bar":
            return Bar
        else:
            return pickle.Unpickler.find_class(self, module, name)

bars = MyUnpickler(open("objects.pkl")).load()

注意注意注意：

如果你是从另一个模块，比如说 baz，调用这段代码，那么解压出来的对象类型将是 baz.Bar，而不是 foo.Bar。假设 foo.Bar 和 baz.Bar 的类定义是一样的，那你解压的时候不会遇到问题。但在后续使用 isinstance、type 等函数时要小心。一般来说，除了偶尔用一次，这种做法可能不太聪明，因为你的代码库里现在有两个 Bar 的实例。如果可能的话，最好把 foo 加入你的路径中。