Ticket #11845: picklefunc.py

r"""
Pickleable function objects

Provides a pickleable wrapper for function objects. The resulting pickles have
the same portability restrictions as marshalled code objects (see the marshal
module in the python standard library)

AUTHORS:

- Nils Bruin (2011-09-24): initial version
"""

import types
import marshal

def makecell(value):
    r"""
    return value wrapped in a cell object
    """
    return (lambda : value).func_closure[0]

class PickleableFunction(object):
    r"""
    Wrapper for functions to make them pickleable
    
    Dynamically defined function objects such as results of lambda expressions
    or from closures, are normally not pickleable. This class provides a
    pickleable wrapper for such function objects. The resulting pickles are only
    portable to the extent that marshalled codeobject are portable (see the
    marshal module in the Python library for a description of the limitations).
    
    INPUT:
    
    - ``f`` - a Python function
    
    OUTPUT:
    
    - A wrapped function that is pickleable
    
    EXAMPLES::
    
        sage: def f(a,b=0):
        ...       "docstring of f"
        ...       return a+b
        sage: p=PickleableFunction(f)
        
    The wrapped object is still callable, but less efficient of course::
    
        sage: p(10,20)
        30
        sage: p
        PickleableFunction(<function f at ...>)
        sage: q=loads(dumps(p))
    
    We cannot expect to hold for reconstructed objects when code is involved:
    
        sage: p == q
        False
        
    But we can throw away the wrapper and call the new object and get the same
    result::
    
        sage: g=q.function
        sage: g(10,20)
        30
        
    Closures of functions are correctly recreated::

        sage: def make_constant(c):
        ....:     return lambda inc=0: c+inc
        sage: f=make_constant(5)
        sage: p=PickleableFunction(f)
        sage: q=loads(dumps(p))
        sage: g=q.function
        sage: g(1)
        6

    Generators are also pickled correctly::
    
        sage: def make_counter(c):
        ....:     def count():
        ....:        for i in xrange(c):
        ....:            yield i
        ....:     return count
        sage: f=make_counter(10)
        sage: g=loads(dumps(PickleableFunction(f))).function
        sage: [i for i in g()]
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
       
    """

    def __init__(self,f):
        """
        """
        if not(isinstance(f,types.FunctionType)):
            raise ErrorType, "PickleableFunction can only wrap python functions"
        self.function=f
            
    def __call__(self,*args,**kwargs):
        """
        Calling the wrapper results in calling the wrapped function
        """
        return self.function(*args,**kwargs)
        
    def __getstate__(self):
        """
        Implementation of pickle protocol
        """
        
        #we extract all data from self.function necessary to recreate it.
        f=self.function

        code=marshal.dumps(f.func_code) #use marshal to serialize code object
        globs=f.func_globals
        name=f.func_name
        kwargs=f.func_defaults
        if f.func_closure:
            # if there is a closure, unwrap the values from their cells
            closure_tuple=tuple(c.cell_contents for c in f.func_closure)
        else:
            closure_tuple=None
        
        #the produced state is a tuple state such that essentially
        #types.FunctionObject(*state)
        #would recreate the right function object. The "code" and "closure"
        #parameters are recoded in a serializable form.
        #we are ignoring the "globals" of the function object because that
        #is almost certainly not meant to be pickled.
        return (code,None,name,kwargs,closure_tuple)
        
    def __setstate__(self,state):
        """
        Implementation of pickle protocol
        """
        
        code,globs,name,kwargs,closure_tuple=state
        code_object=marshal.loads(code) #deserialize code object using marshal
        if globs is None:    #no globals means: use present globals
            globs=globals()
        if closure_tuple:
            #closure needs to be a tuple of cell objects
            closure=tuple(makecell(c) for c in closure_tuple)
        else:
            closure=None
        self.function=types.FunctionType(code_object,globs,name,kwargs,closure)

    def __repr__(self):
        """
        return string representation
        """
        return "PickleableFunction(%s)"%self.function