Ticket #11845: picklefunc.py

File picklefunc.py, 4.5 KB (added by nbruin, 8 years ago)

implementation of a PickleableFunction? class (not a patch!)

Line 
1r"""
2Pickleable function objects
3
4Provides a pickleable wrapper for function objects. The resulting pickles have
5the same portability restrictions as marshalled code objects (see the marshal
6module in the python standard library)
7
8AUTHORS:
9
10- Nils Bruin (2011-09-24): initial version
11"""
12
13import types
14import marshal
15
16def makecell(value):
17    r"""
18    return value wrapped in a cell object
19    """
20    return (lambda : value).func_closure[0]
21
22class PickleableFunction(object):
23    r"""
24    Wrapper for functions to make them pickleable
25   
26    Dynamically defined function objects such as results of lambda expressions
27    or from closures, are normally not pickleable. This class provides a
28    pickleable wrapper for such function objects. The resulting pickles are only
29    portable to the extent that marshalled codeobject are portable (see the
30    marshal module in the Python library for a description of the limitations).
31   
32    INPUT:
33   
34    - ``f`` - a Python function
35   
36    OUTPUT:
37   
38    - A wrapped function that is pickleable
39   
40    EXAMPLES::
41   
42        sage: def f(a,b=0):
43        ...       "docstring of f"
44        ...       return a+b
45        sage: p=PickleableFunction(f)
46       
47    The wrapped object is still callable, but less efficient of course::
48   
49        sage: p(10,20)
50        30
51        sage: p
52        PickleableFunction(<function f at ...>)
53        sage: q=loads(dumps(p))
54   
55    We cannot expect to hold for reconstructed objects when code is involved:
56   
57        sage: p == q
58        False
59       
60    But we can throw away the wrapper and call the new object and get the same
61    result::
62   
63        sage: g=q.function
64        sage: g(10,20)
65        30
66       
67    Closures of functions are correctly recreated::
68
69        sage: def make_constant(c):
70        ....:     return lambda inc=0: c+inc
71        sage: f=make_constant(5)
72        sage: p=PickleableFunction(f)
73        sage: q=loads(dumps(p))
74        sage: g=q.function
75        sage: g(1)
76        6
77
78    Generators are also pickled correctly::
79   
80        sage: def make_counter(c):
81        ....:     def count():
82        ....:        for i in xrange(c):
83        ....:            yield i
84        ....:     return count
85        sage: f=make_counter(10)
86        sage: g=loads(dumps(PickleableFunction(f))).function
87        sage: [i for i in g()]
88        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
89       
90    """
91
92    def __init__(self,f):
93        """
94        """
95        if not(isinstance(f,types.FunctionType)):
96            raise ErrorType, "PickleableFunction can only wrap python functions"
97        self.function=f
98           
99    def __call__(self,*args,**kwargs):
100        """
101        Calling the wrapper results in calling the wrapped function
102        """
103        return self.function(*args,**kwargs)
104       
105    def __getstate__(self):
106        """
107        Implementation of pickle protocol
108        """
109       
110        #we extract all data from self.function necessary to recreate it.
111        f=self.function
112
113        code=marshal.dumps(f.func_code) #use marshal to serialize code object
114        globs=f.func_globals
115        name=f.func_name
116        kwargs=f.func_defaults
117        if f.func_closure:
118            # if there is a closure, unwrap the values from their cells
119            closure_tuple=tuple(c.cell_contents for c in f.func_closure)
120        else:
121            closure_tuple=None
122       
123        #the produced state is a tuple state such that essentially
124        #types.FunctionObject(*state)
125        #would recreate the right function object. The "code" and "closure"
126        #parameters are recoded in a serializable form.
127        #we are ignoring the "globals" of the function object because that
128        #is almost certainly not meant to be pickled.
129        return (code,None,name,kwargs,closure_tuple)
130       
131    def __setstate__(self,state):
132        """
133        Implementation of pickle protocol
134        """
135       
136        code,globs,name,kwargs,closure_tuple=state
137        code_object=marshal.loads(code) #deserialize code object using marshal
138        if globs is None:    #no globals means: use present globals
139            globs=globals()
140        if closure_tuple:
141            #closure needs to be a tuple of cell objects
142            closure=tuple(makecell(c) for c in closure_tuple)
143        else:
144            closure=None
145        self.function=types.FunctionType(code_object,globs,name,kwargs,closure)
146
147    def __repr__(self):
148        """
149        return string representation
150        """
151        return "PickleableFunction(%s)"%self.function