Moving class used in pickle

[Peter Otten]

Jeffrey Barish wrote:

> I have a class derived from string that is used in a pickle.  In the new
> version of my program, I moved the module containing the definition of the
> class.  Now the unpickle fails because it doesn't find the module.  I was
> thinking that I could make the unpickle work by putting a copy of the
> module in the original location and then redefine the class by sticking a
> __setstate__ in the class thusly:
> 
> def __setstate__(self, state):
>     self.__dict__.update(state)
>     self.__class__ = NewClassName
> 
> My plan was to specify the new location of the module in NewClassName.
> However, when I do this I get the message "'class' object layout differs
> from 'class'".  I don't think that they do as the new module is a copy of
> the old one.  I suspect that I am not allowed to make the class assignment
> because my class is derived from string.  What is the best way to update
> the pickle?  The only thought I have is to read all the data with the old
> class module, store the data in some nonpickle format, and then, with
> another program, read the nonpickle-format file and rewrite the pickle
> with the class module in the new location.

You could overwrite Unpickler.find_class():

import pickle
from cStringIO import StringIO

class AliasUnpickler(pickle.Unpickler):
    def __init__(self, aliases, *args, **kw):
        pickle.Unpickler.__init__(self, *args, **kw)
        self.aliases = aliases
    def find_class(self, module, name):
        module, name = self.aliases.get((module, name), (module, name))
        return pickle.Unpickler.find_class(self, module, name)

def loads(aliases, str):
    file = StringIO(str)
    return AliasUnpickler(aliases, file).load()

if __name__ == "__main__":
    import before, after
    data = before.A()
    print data.__class__, data
    dump = pickle.dumps(data)
    data = loads({("before", "A"): ("after", "B")}, dump)
    print data.__class__, data

In the example the aliases dictionary maps (module, classname) pairs to
(module, classname) pairs. Of course this only works when the class layout
wasn't changed.

Peter