diff options
Diffstat (limited to 'intern/python/modules/mcf/utils/cpickle_extend.py')
-rw-r--r-- | intern/python/modules/mcf/utils/cpickle_extend.py | 190 |
1 files changed, 190 insertions, 0 deletions
diff --git a/intern/python/modules/mcf/utils/cpickle_extend.py b/intern/python/modules/mcf/utils/cpickle_extend.py new file mode 100644 index 00000000000..aaca41d51fb --- /dev/null +++ b/intern/python/modules/mcf/utils/cpickle_extend.py @@ -0,0 +1,190 @@ +''' +Extend cpickle storage to include modules, and builtin functions/methods + +To use, just import this module. +''' +import copy_reg + +### OBJECTS WHICH ARE RESTORED THROUGH IMPORTS +# MODULES +def pickle_module(module): + ''' + Store a module to a pickling stream, must be available for + reimport during unpickling + ''' + return unpickle_imported_code, ('import %s'%module.__name__, module.__name__) + +# FUNCTIONS, METHODS (BUILTIN) +def pickle_imported_code(funcmeth): + ''' + Store a reference to an imported element (such as a function/builtin function, + Must be available for reimport during unpickling. + ''' + module = _whichmodule(funcmeth) + return unpickle_imported_code, ('from %s import %s'%(module.__name__,funcmeth.__name__),funcmeth.__name__) + +import types, regex +import_filter = regex.compile('''\(from [A-Za-z0-9_\.]+ \)?import [A-Za-z0-9_\.]+''') # note the limitations on whitespace +getattr_filter = regex.compile('''[A-Za-z0-9_\.]+''') # note we allow you to use x.y.z here + +# MODULES, AND FUNCTIONS +def unpickle_imported_code(impstr,impname): + ''' + Attempt to load a reference to a module or other imported code (such as functions/builtin functions) + ''' + if import_filter.match(impstr) != len(impstr) or getattr_filter.match(impname)!= len(impname): + import sys + sys.stderr.write('''Possible attempt to smuggle arbitrary code into pickle file (see module cpickle_extend).\nPassed code was %s\n%s\n'''%(impstr,impname)) + del(sys) + else: + ns = {} + try: + exec (impstr) in ns # could raise all sorts of errors, of course, and is still dangerous when you have no control over the modules on your system! Do not allow for untrusted code!!! + return eval(impname, ns) + except: + import sys + sys.stderr.write('''Error unpickling module %s\n None returned, will likely raise errors.'''%impstr) + return None + +# Modules +copy_reg.pickle(type(regex),pickle_module,unpickle_imported_code) +# builtin functions/methods +copy_reg.pickle(type(regex.compile),pickle_imported_code, unpickle_imported_code) + +del(regex) # to keep the namespace neat as possible + +### INSTANCE METHODS +''' +The problem with instance methods is that they are almost always +stored inside a class somewhere. We really need a new type: reference +that lets us just say "y.this" + +We also need something that can reliably find burried functions :( not +likely to be easy or clean... + +then filter for x is part of the set +''' +import new + +def pickle_instance_method(imeth): + ''' + Use the (rather surprisingly clean) internals of + the method to store a reference to a method. Might + be better to use a more general "get the attribute + 'x' of this object" system, but I haven't written that yet :) + ''' + klass = imeth.im_class + funcimp = _imp_meth(imeth) + self = imeth.im_self # will be None for UnboundMethodType + return unpickle_instance_method, (funcimp,self,klass) +def unpickle_instance_method(funcimp,self,klass): + ''' + Attempt to restore a reference to an instance method, + the instance has already been recreated by the system + as self, so we just call new.instancemethod + ''' + funcimp = apply(unpickle_imported_code, funcimp) + return new.instancemethod(func,self,klass) + +copy_reg.pickle(types.MethodType, pickle_instance_method, unpickle_instance_method) +copy_reg.pickle(types.UnboundMethodType, pickle_instance_method, unpickle_instance_method) + +### Arrays +try: + import array + LittleEndian = array.array('i',[1]).tostring()[0] == '\001' + def pickle_array(somearray): + ''' + Store a standard array object, inefficient because of copying to string + ''' + return unpickle_array, (somearray.typecode, somearray.tostring(), LittleEndian) + def unpickle_array(typecode, stringrep, origendian): + ''' + Restore a standard array object + ''' + newarray = array.array(typecode) + newarray.fromstring(stringrep) + # floats are always big-endian, single byte elements don't need swapping + if origendian != LittleEndian and typecode in ('I','i','h','H'): + newarray.byteswap() + return newarray + copy_reg.pickle(array.ArrayType, pickle_array, unpickle_array) +except ImportError: # no arrays + pass + +### NUMPY Arrays +try: + import Numeric + LittleEndian = Numeric.array([1],'i').tostring()[0] == '\001' + def pickle_numpyarray(somearray): + ''' + Store a numpy array, inefficent, but should work with cPickle + ''' + return unpickle_numpyarray, (somearray.typecode(), somearray.shape, somearray.tostring(), LittleEndian) + def unpickle_numpyarray(typecode, shape, stringval, origendian): + ''' + Restore a numpy array + ''' + newarray = Numeric.fromstring(stringval, typecode) + Numeric.reshape(newarray, shape) + if origendian != LittleEndian and typecode in ('I','i','h','H'): + # this doesn't seem to work correctly, what's byteswapped doing??? + return newarray.byteswapped() + else: + return newarray + copy_reg.pickle(Numeric.ArrayType, pickle_numpyarray, unpickle_numpyarray) +except ImportError: + pass + +### UTILITY FUNCTIONS +classmap = {} +def _whichmodule(cls): + """Figure out the module in which an imported_code object occurs. + Search sys.modules for the module. + Cache in classmap. + Return a module name. + If the class cannot be found, return __main__. + Copied here from the standard pickle distribution + to prevent another import + """ + if classmap.has_key(cls): + return classmap[cls] + clsname = cls.__name__ + for name, module in sys.modules.items(): + if name != '__main__' and \ + hasattr(module, clsname) and \ + getattr(module, clsname) is cls: + break + else: + name = '__main__' + classmap[cls] = name + return name + +import os, string, sys + +def _imp_meth(im): + ''' + One-level deep recursion on finding methods, i.e. we can + find them only if the class is at the top level. + ''' + fname = im.im_func.func_code.co_filename + tail = os.path.splitext(os.path.split(fname)[1])[0] + ourkeys = sys.modules.keys() + possibles = filter(lambda x,tail=tail: x[-1] == tail, map(string.split, ourkeys, ['.']*len(ourkeys))) + + # now, iterate through possibles to find the correct class/function + possibles = map(string.join, possibles, ['.']*len(possibles)) + imp_string = _search_modules(possibles, im.im_func) + return imp_string + +def _search_modules(possibles, im_func): + for our_mod_name in possibles: + our_mod = sys.modules[our_mod_name] + if hasattr(our_mod, im_func.__name__) and getattr(our_mod, im_func.__name__).im_func is im_func: + return 'from %s import %s'%(our_mod.__name__, im_func.__name__), im_func.__name__ + for key,val in our_mod.__dict__.items(): + if hasattr(val, im_func.__name__) and getattr(val, im_func.__name__).im_func is im_func: + return 'from %s import %s'%(our_mod.__name__,key), '%s.%s'%(key,im_func.__name__) + raise '''No import string calculable for %s'''%im_func + + |