How to substitute module.Class() to locally defined Class() when loading with Python's Pickle?

Question

Hi,

I have a pickle dump which has an array of foo.Bar() objects. I am trying to unpickle it, but the Bar() class definition is in the same file that's trying to unpickle, and not in the foo module. So, pickle complains that it couldn't find module foo.

I tried to inject foo module doing something similar to: import imp, sys

class Bar:
    pass

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

import foo
print foo.Bar()

Which works, but when I try to add, after that:

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

I receive the friendly error:

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

Any ideas?

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()

CAVEAT CAVEAT CAVEAT: If you are calling this snippet of code from another module, let's say baz, then the unpickled objects will be of type baz.Bar, not foo.Bar. Assuming the class definitions of foo.Bar and baz.Bar are the same, you will have no trouble unpickling. But be careful with downstream use of isinstance, type, etc. In general, trying to do this sort of thing for anything but a one-off is probably not smart, because your codebase now contains two instances of Bar. If at all possible you should just put foo in your PATH.