Merge pull request #1708 from robertwb/pickle

Support pickling of cdef classes.

CHANGES.rst

@@ -22,6 +22,8 @@ Features added
 * ``abs()`` is optimised for C complex numbers.
   Patch by da-woods (Github issue #1648).
 
+* cdef classes now support pickling by default when possible.
+
 Bugs fixed
 ----------
 

Cython/Compiler/Options.py

@@ -144,6 +144,7 @@ _directive_defaults = {
     'embedsignature' : False,
     'locals' : {},
     'auto_cpdef': False,
+    'auto_pickle': None,
     'cdivision': False, # was True before 0.12
     'cdivision_warnings': False,
     'overflowcheck': False,
@@ -263,6 +264,7 @@ def normalise_encoding_name(option_name, encoding):
 
 # Override types possibilities above, if needed
 directive_types = {
+    'auto_pickle': bool,
     'final' : bool,  # final cdef classes and methods
     'internal' : bool,  # cdef class visibility in the module dict
     'infer_types' : bool, # values can be True/None/False
@@ -285,6 +287,7 @@ for key, val in _directive_defaults.items():
 
 directive_scopes = { # defaults to available everywhere
     # 'module', 'function', 'class', 'with statement'
+    'auto_pickle': ('module', 'cclass'),
     'final' : ('cclass', 'function'),
     'inline' : ('function',),
     'staticmethod' : ('function',),  # FIXME: analysis currently lacks more specific function scope

Cython/Compiler/ParseTreeTransforms.py

@@ -1535,10 +1535,13 @@ if VALUE is not None:
         return node
 
     def visit_ModuleNode(self, node):
+        # Pickling support requires injecting module-level nodes.
+        self.extra_module_declarations = []
         self.seen_vars_stack.append(set())
         node.analyse_declarations(self.current_env())
         self.visitchildren(node)
         self.seen_vars_stack.pop()
+        node.body.stats.extend(self.extra_module_declarations)
         return node
 
     def visit_LambdaNode(self, node):
@@ -1560,8 +1563,93 @@ if VALUE is not None:
                     stats.append(property)
             if stats:
                 node.body.stats += stats
+            if (node.visibility != 'extern'
+                and 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):
+        env = self.current_env()
+        if node.scope.directives['auto_pickle'] is False:   # None means attempt it.
+            # Old behavior of not doing anything.
+            return
+
+        all_members = []
+        cls = node.entry.type
+        cinit = None
+        while cls is not None:
+            all_members.extend(e for e in cls.scope.var_entries if e.name not in ('__weakref__', '__dict__'))
+            cinit = cinit or cls.scope.lookup('__cinit__')
+            cls = cls.base_type
+        all_members.sort(key=lambda e: e.name)
+
+        non_py = [
+            e for e in all_members
+            if not e.type.is_pyobject and (not e.type.create_from_py_utility_code(env)
+                                           or not e.type.create_to_py_utility_code(env))]
+
+        if cinit or non_py:
+            if cinit:
+                # TODO(robertwb): We could allow this if __cinit__ has no require arguments.
+                msg = 'no default __reduce__ due to non-trivial __cinit__'
+            else:
+                msg = "%s cannot be converted to a Python object for pickling" % ','.join("self.%s" % e.name for e in non_py)
+
+            if node.scope.directives['auto_pickle'] is True:
+                error(node.pos, msg)
+
+            pickle_func = TreeFragment(u"""
+                def __reduce__(self):
+                    raise TypeError("%s")
+                """ % msg,
+                level='c_class', pipeline=[NormalizeTree(None)]).substitute({})
+            pickle_func.analyse_declarations(node.scope)
+            self.visit(pickle_func)
+            node.body.stats.append(pickle_func)
+
+        else:
+            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(__pyx_type, __pyx_state, %(args)s):
+                    cdef %(class_name)s result
+                    result = %(class_name)s.__new__(__pyx_type)
+                    %(assignments)s
+                    if hasattr(result, '__setstate__'):
+                        result.__setstate__(__pyx_state)
+                    elif hasattr(result, '__dict__'):
+                        result.__dict__.update(__pyx_state)
+                    elif __pyx_state is not None:
+                        from pickle import PickleError
+                        raise PickleError("Unexpected state: %%s" %% __pyx_state)
+                    return result
+                """ % {
+                    'unpickle_func_name': unpickle_func_name,
+                    'class_name': node.class_name,
+                    'assignments': '; '.join('result.%s = __pyx_%s' % (v, v) for v in all_members_names),
+                    'args': ','.join('__pyx_%s' % v for v in 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):
+                    if hasattr(self, '__getstate__'):
+                        state = self.__getstate__()
+                    elif hasattr(self, '__dict__'):
+                        state = self.__dict__
+                    else:
+                        state = None
+                    return %s, (type(self), state, %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):
         """
         Create function calls to the decorators and reassignments to

docs/src/userguide/extension_types.rst

@@ -550,6 +550,23 @@ If you can be sure addresses will contain only references to strings,
 the above would be safe, and it may yield a significant speedup, depending on
 your usage pattern.
 
+
+Controlling pickling
+====================
+
+By default, Python will generate a ``__reduce__`` to allow pickling an extension
+type if and only if each of its members are convertible to Python and it has
+no ``__cinit__`` method.
+To require this behavior (i.e. throw an error at compile time if a class
+cannot be pickled) decorate the class with ``@cython.auto_pickle(True)``.
+One can also annotate with ``@cython.auto_pickle(False)`` to get the old
+behavior of not generating a ``__reduce__`` method in any case.
+
+Manually implementing a ``__reduce__`` or `__reduce_ex__`` method will also
+disable this auto-generation and can be used to support pickling of more
+complicated types.
+
+
 Public and external extension types
 ====================================
 

tests/run/reduce_pickle.pyx

+import cython
 import sys
 
 if sys.version_info[0] < 3:
@@ -57,3 +58,81 @@ cdef class B:
 
 def makeB(kwds):
     return B(**kwds)
+
+
+@cython.auto_pickle(True)  # Not needed, just to test the directive.
+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 readonly int i
+    cdef readonly 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)
+
+
+class DefaultReducePySubclass(DefaultReduce):
+    """
+    >>> a = DefaultReducePySubclass(i=11, s='abc', x=1.5); a
+    DefaultReducePySubclass(i=11, s='abc', x=1.5)
+    >>> import pickle
+    >>> pickle.loads(pickle.dumps(a))
+    DefaultReducePySubclass(i=11, s='abc', x=1.5)
+    """
+    def __init__(self, **kwargs):
+        self.x = kwargs.pop('x', 0)
+        super(DefaultReducePySubclass, self).__init__(**kwargs)
+
+    def __repr__(self):
+        return "DefaultReducePySubclass(i=%s, s=%r, x=%s)" % (self.i, self.s, self.x)
+
+
+cdef class NoReduceDueToIntPtr(object):
+    """
+    >>> import pickle
+    >>> pickle.dumps(NoReduceDueToIntPtr())
+    Traceback (most recent call last):
+    ...
+    TypeError: self.int_ptr cannot be converted to a Python object for pickling
+    """
+    cdef int* int_ptr
+
+cdef class NoReduceDueToNontrivialCInit(object):
+    """
+    >>> import pickle
+    >>> pickle.dumps(NoReduceDueToNontrivialCInit(None))
+    Traceback (most recent call last):
+    ...
+    TypeError: no default __reduce__ due to non-trivial __cinit__
+    """
+    def __cinit__(self, arg):
+        pass