Source Code for Module topo.misc.picklemain

  1  """ 
  2  Extensions to pickle allowing items in __main__ to be saved. 
  3  """ 
  4   
  5  import new 
  6  import pickle 
  7  import new 
  8  import types 
  9  import __main__ 
 10  from StringIO import StringIO 
 11   
 12   
 13   
 14 -class PickleMain(object): 
 15      """ 
 16      Pickle support for types and functions defined in __main__. 
 17       
 18      When pickled, saves types and functions defined in __main__ by 
 19      value (i.e. as bytecode). When unpickled, loads previously saved 
 20      types and functions back into __main__. 
 21      """ 
 22 -    def _create_pickler(self): 
 23          # Usually we use the cPickle module rather than the pickle 
 24          # module (because cPickle is faster), but here we need control 
 25          # over the pickling process so we use pickle. 
 26          # 
 27          # Additionally, we create a Pickler instance to avoid changing 
 28          # defaults in the pickle module itself, so that there are no side 
 29          # effects for code elsewhere (although we don't use pickle 
 30          # anywhere else ourselves...). 
 31   
 32          self.pickled_bytecode = StringIO() 
 33          self.pickler = pickle.Pickler(self.pickled_bytecode,-1) 
 34   
 35          self.pickler.dispatch[new.code] = save_code 
 36          self.pickler.dispatch[new.function] = save_function 
 37          self.pickler.dispatch[dict] = save_module_dict 
 38          self.pickler.dispatch[new.classobj] = save_classobj 
 39          self.pickler.dispatch[new.instancemethod] = save_instancemethod 
 40          self.pickler.dispatch[new.module] = save_module 
 41          self.pickler.dispatch[type] = save_type 
 42          # CB: maybe this should be registered from elsewhere 
 43          import param 
 44          self.pickler.dispatch[param.parameterized.ParameterizedMetaclass] = save_type 
 45   
 46   
 47 -    def __getstate__(self): 
 48          self._create_pickler() 
 49           
 50          bytecode = {} 
 51          for name,obj in __main__.__dict__.items(): 
 52              if not name.startswith('_'): 
 53                  if isinstance(obj,types.FunctionType) or isinstance(obj,type): 
 54                      # (could be extended to other types, I guess 
 55                      if obj.__module__ == "__main__": 
 56                          #CB: how do I print out info via Parameterized? 
 57                          print "%s is defined in __main__: saving bytecode."%name 
 58                          bytecode[name] = obj 
 59   
 60          self.pickler.dump(bytecode) 
 61          return {'pickled_bytecode':self.pickled_bytecode} 
 62   
 63   
 64 -    def __setstate__(self,state): 
 65          bytecode = pickle.load(StringIO(state['pickled_bytecode'].getvalue())) 
 66   
 67          for name,obj in bytecode.items(): 
 68              print "%s restored from bytecode into __main__"%name 
 69              __main__.__dict__[name] = obj 
 70   
 71   
 72   
 73   
 74   
 75  ### Copied from http://code.activestate.com/recipes/572213/ ### 
 76   
 77  # Original docstring 
 78  # 
 79  # Extend pickle module to allow pickling of interpreter state 
 80  # including any interactively defined functions and classes. 
 81  # 
 82  # This module is not required for unpickling such pickle files. 
 83  # 
 84  # >>> import savestate, pickle, __main__ 
 85  # >>> pickle.dump(__main__, open('savestate.pickle', 'wb'), 2) 
 86   
 87  # (note that we're not actually using it to pickle __main__, and I've 
 88  # removed the lines that change pickle's defaults) 
 89   
 90   
 91 -def save_code(self, obj): 
 92      """ Save a code object by value """ 
 93      args = ( 
 94          obj.co_argcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, 
 95          obj.co_consts, obj.co_names, obj.co_varnames, obj.co_filename, obj.co_name, 
 96          obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, obj.co_cellvars 
 97      ) 
 98      self.save_reduce(new.code, args, obj=obj) 
 99   
100 -def save_function(self, obj): 
101      """ Save functions by value if they are defined interactively """ 
102      if obj.__module__ == '__main__' or obj.func_name == '<lambda>': 
103          args = (obj.func_code, obj.func_globals, obj.func_name, obj.func_defaults, obj.func_closure) 
104          self.save_reduce(new.function, args, obj=obj) 
105      else: 
106          pickle.Pickler.save_global(self, obj) 
107   
108 -def save_global_byname(self, obj, modname, objname): 
109      """ Save obj as a global reference. Used for objects that pickle does not find correctly. """ 
110      self.write('%s%s\n%s\n' % (pickle.GLOBAL, modname, objname)) 
111      self.memoize(obj) 
112   
113 -def save_module_dict(self, obj, main_dict=vars(__import__('__main__'))): 
114      """ Special-case __main__.__dict__. Useful for a function's func_globals member. """ 
115      if obj is main_dict: 
116          save_global_byname(self, obj, '__main__', '__dict__') 
117      else: 
118          return pickle.Pickler.save_dict(self, obj)      # fallback to original  
119   
120 -def save_classobj(self, obj): 
121      """ Save an interactively defined classic class object by value """ 
122      if obj.__module__ == '__main__': 
123          args = (obj.__name__, obj.__bases__, obj.__dict__) 
124          self.save_reduce(new.classobj, args, obj=obj) 
125      else:  # CEBALERT: clearly wrong ('name' not defined). This code 
126             # never activated? 
127          pickle.Pickler.save_global(self, obj, name) 
128   
129 -def save_instancemethod(self, obj): 
130      """ Save an instancemethod object """ 
131      # Instancemethods are re-created each time they are accessed so this will not be memoized 
132      args = (obj.im_func, obj.im_self, obj.im_class) 
133      self.save_reduce(new.instancemethod, args) 
134   
135 -def save_module(self, obj): 
136      """ Save modules by reference, except __main__ which also gets its contents saved by value """ 
137      if obj.__name__ == '__main__': 
138          self.save_reduce(__import__, (obj.__name__,), obj=obj, state=vars(obj).copy()) 
139      elif obj.__name__.count('.') == 0: 
140          self.save_reduce(__import__, (obj.__name__,), obj=obj) 
141      else: 
142          save_global_byname(self, obj, *obj.__name__.rsplit('.', 1)) 
143   
144 -def save_type(self, obj): 
145      if getattr(new, obj.__name__, None) is obj: 
146          # Types in 'new' module claim their module is '__builtin__' but are not actually there 
147          save_global_byname(self, obj, 'new', obj.__name__) 
148      elif obj.__module__ == '__main__': 
149          # Types in __main__ are saved by value 
150   
151          # Make sure we have a reference to type.__new__         
152          if id(type.__new__) not in self.memo: 
153              self.save_reduce(getattr, (type, '__new__'), obj=type.__new__) 
154              self.write(pickle.POP) 
155   
156          # Copy dictproxy to real dict 
157          d = dict(obj.__dict__) 
158          # Clean up unpickleable descriptors added by Python 
159          d.pop('__dict__', None) 
160          d.pop('__weakref__', None) 
161           
162          args = (type(obj), obj.__name__, obj.__bases__, d) 
163          self.save_reduce(type.__new__, args, obj=obj) 
164      else: 
165          # Fallback to default behavior: save by reference 
166          pickle.Pickler.save_global(self, obj) 
167  ############################################################### 
168