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
+
+