| 1 | import pickle | 
|---|
| 2 | from cStringIO import StringIO | 
|---|
| 3 |  | 
|---|
| 4 | class 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 |  | 
|---|
| 12 | class 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() | 
|---|