Stackless issue python#284: Update pickling of functions

Update the pickling code to serialize all attributes of function
objects.

Author: akruis

Stackless/changelog.txt

@@ -9,6 +9,12 @@ What's New in Stackless 3.X.X?
 
 *Release date: 20XX-XX-XX*
 
+- https://github.com/stackless-dev/stackless/issues/284
+  Stackless supports pickling functions. But previous versions of Stackless
+  only pickled the following attributes of a function-object: '__code__',
+  '__globals__', '__name__', '__defaults__', '__closure__' and '__module__'.
+  Now Stackless pickles all attributes of a function.
+
 - https://github.com/stackless-dev/stackless/issues/268
   Pickled code objects now contain importlib.util.MAGIC_NUMBER. Unpickling
   code objects pickled with an incompatible version of Python now creates a

Stackless/pickling/prickelpit.c

@@ -820,25 +820,42 @@ static int init_celltype(PyObject * mod)
 
  ******************************************************/
 
-#define functuplefmt "OOOOOO"
+#define functuplefmt_pre38 "OOOOOO"
+#define functuplefmt functuplefmt_pre38 "OOOOO"
 
 static PyTypeObject wrap_PyFunction_Type;
 
 static PyObject *
 func_reduce(PyFunctionObject * func, PyObject *unused)
 {
+    /* See funcobject.c: some attribute can't be NULL. */
+    assert(func->func_code);
+    assert(func->func_globals);
+    assert(func->func_name);
+    assert(func->func_doc);
+    assert(func->func_qualname);
+
+    PyObject *dict = PyObject_GenericGetDict((PyObject *)func, NULL);
+    if (NULL == dict)
+        return NULL;
+
     PyObject *tup = Py_BuildValue(
         "(O()(" functuplefmt "))",
         &wrap_PyFunction_Type,
         /* Standard function constructor arguments. */
-        func->func_code != NULL ? func->func_code : Py_None,
-        func->func_globals != NULL ? func->func_globals : Py_None,
-        func->func_name != NULL ? func->func_name : Py_None,
+        func->func_code,
+        func->func_globals,
+        func->func_name,
         func->func_defaults != NULL ? func->func_defaults : Py_None,
         func->func_closure != NULL ? func->func_closure : Py_None,
-        /* Additional data we need to preserve. */
-        func->func_module != NULL ? func->func_module : Py_None
+        func->func_module != NULL ? func->func_module : Py_None,
+        func->func_kwdefaults != NULL ? func->func_kwdefaults : Py_None,
+        func->func_doc,
+        dict,
+        func->func_annotations != NULL ? func->func_annotations : Py_None,
+        func->func_qualname
     );
+    Py_DECREF(dict);
     return tup;
 }
 
@@ -893,25 +910,58 @@ func_setstate(PyObject *self, PyObject *args)
     if (args2 == NULL)
         return NULL;
 
-    fu = (PyFunctionObject *)
-        Py_TYPE(self)->tp_new(Py_TYPE(self), args2, NULL);
+    fu = (PyFunctionObject *) Py_TYPE(self)->tp_new(Py_TYPE(self), args2, NULL);
     Py_DECREF(args2);
-    if (fu != NULL) {
-        PyFunctionObject *target = (PyFunctionObject *) self;
-        COPY(fu, target, func_code);
-        COPY(fu, target, func_globals);
-        COPY(fu, target, func_name);
-        COPY(fu, target, func_defaults);
-        COPY(fu, target, func_closure);
+    if (fu == NULL)
+        return NULL;
+
+    PyFunctionObject *target = (PyFunctionObject *) self;
+    COPY(fu, target, func_code);
+    COPY(fu, target, func_globals);
+    COPY(fu, target, func_name);
+    COPY(fu, target, func_defaults);
+    COPY(fu, target, func_closure);
+    Py_DECREF(fu);
+
+    args2 = PyTuple_GetItem(args, 5);
+    if (NULL == args2)
+        return NULL;
+    Py_INCREF(args2);
+    Py_XSETREF(target->func_module, args2);
+
+    if (PyTuple_GET_SIZE(args) != sizeof(functuplefmt_pre38)-1) {
+        /* Stackless 3.8 and up */
+        if (PyTuple_GET_SIZE(args) != sizeof(functuplefmt)-1) {
+            PyErr_Format(PyExc_IndexError, "function.__setstate__ expects a tuple of length %d", (int)sizeof(functuplefmt)-1);
+            return NULL;
+        }
+        args2 = PyTuple_GET_ITEM(args, 6);
+        if (PyFunction_SetKwDefaults(self, args2))
+            return NULL;
 
-        Py_XINCREF(PyTuple_GetItem(args, 5));
-    target->func_module = PyTuple_GetItem(args, 5);
+        args2 = PyTuple_GET_ITEM(args, 7);
+        Py_INCREF(args2);
+        Py_XSETREF(target->func_doc, args2);
 
-        Py_DECREF(fu);
-        Py_INCREF(self);
-        return self;
+        args2 = PyTuple_GET_ITEM(args, 8);
+        if (args2 != Py_None && PyObject_GenericSetDict(self, args2, NULL))
+            return NULL;
+
+        args2 = PyTuple_GET_ITEM(args, 9);
+        if (PyFunction_SetAnnotations(self, args2))
+            return NULL;
+
+        args2 = PyTuple_GET_ITEM(args, 10);
+        if(!PyUnicode_Check(args2)) {
+            PyErr_SetString(PyExc_TypeError, "__qualname__ must be set to a string object");
+            return NULL;
+        }
+        Py_INCREF(args2);
+        Py_XSETREF(target->func_qualname, args2);
     }
-    return NULL;
+
+    Py_INCREF(self);
+    return self;
 }
 
 #undef COPY

Stackless/unittests/test_pickle.py

@@ -499,12 +499,6 @@ def testSubmoduleImporting(self):
         m2 = self.loads(self.dumps(m1))
         self.assertEqual(m1, m2)
 
-    def testFunctionModulePreservation(self):
-        # The 'module' name on the function was not being preserved.
-        f1 = lambda: None  # @IgnorePep8
-        f2 = self.loads(self.dumps(f1))
-        self.assertEqual(f1.__module__, f2.__module__)
-
 
 class TestFramePickling(StacklessTestCase):
     def test_get_set_reduce_frame(self):
@@ -1215,9 +1209,14 @@ def create_cell(obj):
         self._test(obj, 'cell_contents')
 
     def test_function(self):
-        def obj():
-            pass
-        self._test(obj, '__code__', '__globals__', '__name__', '__defaults__', '__closure__')
+        v = 123
+        def obj(a, b, c=4711, *, d=815):
+            "a test function"
+            return v
+        obj.__annotations__ = {'bla': 'fasel'}
+        self._test(obj, '__code__', '__globals__', '__name__', '__defaults__',
+                   '__closure__', '__module__', '__kwdefaults__', '__doc__',
+                   '__dict__', '__annotations__', '__qualname__')
 
     def test_frame(self):
         obj = sys._getframe()
@@ -1461,6 +1460,30 @@ def test_run_with_wrong_magic(self):
         self.assertEqual(rc, 42)
 
 
+class TestFunctionPickling(StacklessPickleTestCase):
+    def setUp(self):
+        super().setUp()
+
+    def tearDown(self):
+        super().tearDown()
+
+    def testFunctionModulePreservation(self):
+        # The 'module' name on the function was not being preserved.
+        f1 = lambda: None  # @IgnorePep8
+        f2 = self.loads(self.dumps(f1))
+        self.assertEqual(f1.__module__, f2.__module__)
+
+    def test_unpickle_pre38(self):
+        def obj():
+            pass
+
+        reduced = stackless._stackless._wrap.function.__reduce__(obj)
+        obj2 = reduced[0](*reduced[1])
+        args = reduced[2][:6]
+        obj2.__setstate__(args)
+        self.assertIs(type(obj2), type(obj))
+
+
 if __name__ == '__main__':
     if not sys.argv[1:]:
         sys.argv.append('-v')