Commit 437fcbe8 authored by Robert Bradshaw's avatar Robert Bradshaw

Basic pickling suport.

parent 296e07f4
...@@ -1535,10 +1535,13 @@ if VALUE is not None: ...@@ -1535,10 +1535,13 @@ if VALUE is not None:
return node return node
def visit_ModuleNode(self, node): def visit_ModuleNode(self, node):
# Pickling support requires injecting module-level nodes.
self.extra_module_declarations = []
self.seen_vars_stack.append(set()) self.seen_vars_stack.append(set())
node.analyse_declarations(self.current_env()) node.analyse_declarations(self.current_env())
self.visitchildren(node) self.visitchildren(node)
self.seen_vars_stack.pop() self.seen_vars_stack.pop()
node.body.stats.extend(self.extra_module_declarations)
return node return node
def visit_LambdaNode(self, node): def visit_LambdaNode(self, node):
...@@ -1560,7 +1563,44 @@ if VALUE is not None: ...@@ -1560,7 +1563,44 @@ if VALUE is not None:
stats.append(property) stats.append(property)
if stats: if stats:
node.body.stats += stats node.body.stats += stats
return node if not node.scope.lookup('__reduce__') and not node.scope.lookup('__reduce_ex__'):
self._inject_pickle_methods(node)
return node
def _inject_pickle_methods(self, node):
all_members = []
cls = node.entry.type
while cls is not None:
all_members.extend(cls.scope.var_entries)
cls = cls.base_type
all_members.sort(key=lambda e: e.name)
all_members_names = [e.name for e in all_members]
unpickle_func_name = '__pyx_unpickle_%s' % node.class_name
unpickle_func = TreeFragment(u"""
def %(unpickle_func_name)s(%(args)s):
cdef %(class_name)s result
result = %(class_name)s.__new__(%(class_name)s)
%(assignments)s
return result
""" % {
'unpickle_func_name': unpickle_func_name,
'class_name': node.class_name,
'assignments': '; '.join('result.%s = %s' % (v, v) for v in all_members_names),
'args': ','.join(all_members_names),
}, level='module', pipeline=[NormalizeTree(None)]).substitute({})
unpickle_func.analyse_declarations(node.entry.scope)
self.visit(unpickle_func)
self.extra_module_declarations.append(unpickle_func)
pickle_func = TreeFragment(u"""
def __reduce__(self):
return %s, (%s)
""" % (unpickle_func_name, ', '.join('self.%s' % v for v in all_members_names)),
level='c_class', pipeline=[NormalizeTree(None)]).substitute({})
pickle_func.analyse_declarations(node.scope)
self.visit(pickle_func)
node.body.stats.append(pickle_func)
def _handle_fused_def_decorators(self, old_decorators, env, node): def _handle_fused_def_decorators(self, old_decorators, env, node):
""" """
......
...@@ -57,3 +57,42 @@ cdef class B: ...@@ -57,3 +57,42 @@ cdef class B:
def makeB(kwds): def makeB(kwds):
return B(**kwds) return B(**kwds)
cdef class DefaultReduce(object):
"""
>>> a = DefaultReduce(11, 'abc'); a
DefaultReduce(i=11, s='abc')
>>> import pickle
>>> pickle.loads(pickle.dumps(a))
DefaultReduce(i=11, s='abc')
"""
cdef int i
cdef str s
def __init__(self, i=0, s=None):
self.i = i
self.s = s
def __repr__(self):
return "DefaultReduce(i=%s, s=%r)" % (self.i, self.s)
cdef class DefaultReduceSubclass(DefaultReduce):
"""
>>> a = DefaultReduceSubclass(i=11, s='abc', x=1.5); a
DefaultReduceSubclass(i=11, s='abc', x=1.5)
>>> import pickle
>>> pickle.loads(pickle.dumps(a))
DefaultReduceSubclass(i=11, s='abc', x=1.5)
"""
cdef double x
def __init__(self, **kwargs):
self.x = kwargs.pop('x', 0)
super(DefaultReduceSubclass, self).__init__(**kwargs)
def __repr__(self):
return "DefaultReduceSubclass(i=%s, s=%r, x=%s)" % (self.i, self.s, self.x)
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment