/lib/galaxy/util/aliaspickler.py

 1import pickle
 2from cStringIO import StringIO
 3
 4class AliasUnpickler( pickle.Unpickler ):
 5    def __init__( self, aliases, *args, **kw):
 6        pickle.Unpickler.__init__( self, *args, **kw )
 7        self.aliases = aliases
 8    def find_class( self, module, name ):
 9        module, name = self.aliases.get((module,name), (module,name))
10        return pickle.Unpickler.find_class( self, module, name )
11
12class AliasPickleModule( object ):
13    def __init__( self, aliases ):
14        self.aliases = aliases
15    def dump( self, obj, fileobj, protocol=0):
16        return pickle.dump( obj, fileobj, protocol )
17    def dumps( self, obj, protocol=0 ):
18        return pickle.dumps( obj, protocol )
19    def load( self, fileobj ):
20        return AliasUnpickler( self.aliases, fileobj ).load()
21    def loads( self, string ):
22        fileobj = StringIO( string )
23        return AliasUnpickler( self.aliases, fileobj ).load()