""" Extend pickle module to allow pickling of interpreter state including any interactively defined functions and classes. This module is not required for unpickling such pickle files. >>> import savestate, pickle, __main__ >>> pickle.dump(__main__, open('savestate.pickle', 'wb'), 2) """ import sys, pickle, new def save_code(self, obj): """ Save a code object by value """ args = ( obj.co_argcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, obj.co_varnames, obj.co_filename, obj.co_name, obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, obj.co_cellvars ) self.save_reduce(new.code, args, obj=obj) pickle.Pickler.dispatch[new.code] = save_code def save_function(self, obj): """ Save functions by value if they are defined interactively """ if obj.__module__ == '__main__' or obj.func_name == '': args = (obj.func_code, obj.func_globals, obj.func_name, obj.func_defaults, obj.func_closure) self.save_reduce(new.function, args, obj=obj) else: pickle.Pickler.save_global(self, obj) pickle.Pickler.dispatch[new.function] = save_function def save_global_byname(self, obj, modname, objname): """ Save obj as a global reference. Used for objects that pickle does not find correctly. """ self.write('%s%s\n%s\n' % (pickle.GLOBAL, modname, objname)) self.memoize(obj) def save_module_dict(self, obj, main_dict=vars(__import__('__main__'))): """ Special-case __main__.__dict__. Useful for a function's func_globals member. """ if obj is main_dict: save_global_byname(self, obj, '__main__', '__dict__') else: return pickle.Pickler.save_dict(self, obj) # fallback to original pickle.Pickler.dispatch[dict] = save_module_dict def save_classobj(self, obj): """ Save an interactively defined classic class object by value """ if obj.__module__ == '__main__': args = (obj.__name__, obj.__bases__, obj.__dict__) self.save_reduce(new.classobj, args, obj=obj) else: pickle.Pickler.save_global(self, obj, name) pickle.Pickler.dispatch[new.classobj] = save_classobj def save_instancemethod(self, obj): """ Save an instancemethod object """ # Instancemethods are re-created each time they are accessed so this will not be memoized args = (obj.im_func, obj.im_self, obj.im_class) self.save_reduce(new.instancemethod, args) pickle.Pickler.dispatch[new.instancemethod] = save_instancemethod def save_module(self, obj): """ Save modules by reference, except __main__ which also gets its contents saved by value """ if obj.__name__ == '__main__': self.save_reduce(__import__, (obj.__name__,), obj=obj, state=vars(obj).copy()) elif obj.__name__.count('.') == 0: self.save_reduce(__import__, (obj.__name__,), obj=obj) else: save_global_byname(self, obj, *obj.__name__.rsplit('.', 1)) pickle.Pickler.dispatch[new.module] = save_module def save_type(self, obj): if getattr(new, obj.__name__, None) is obj: # Types in 'new' module claim their module is '__builtin__' but are not actually there save_global_byname(self, obj, 'new', obj.__name__) elif obj.__module__ == '__main__': # Types in __main__ are saved by value # Make sure we have a reference to type.__new__ if id(type.__new__) not in self.memo: self.save_reduce(getattr, (type, '__new__'), obj=type.__new__) self.write(pickle.POP) # Copy dictproxy to real dict d = dict(obj.__dict__) # Clean up unpickleable descriptors added by Python d.pop('__dict__', None) d.pop('__weakref__', None) args = (type(obj), obj.__name__, obj.__bases__, d) self.save_reduce(type.__new__, args, obj=obj) else: # Fallback to default behavior: save by reference pickle.Pickler.save_global(self, obj) pickle.Pickler.dispatch[type] = save_type