Issue python#107: disable pickling of frames

This is part 2 of issue python#107. Second step. Stackless no longer pretends to be able to pickle frame objects in general. It still pickles frames as part of tasklets or traceback objects. This way Stackless adheres closely to the pickling protocol.

https://bitbucket.org/stackless-dev/stackless/issues/107

Author: Anselm Kruis

Lib/stackless.py

@@ -79,9 +79,48 @@ def transmogrify():
     from copyreg import pickle
     for name in dir(_wrap):
         cls = getattr(_wrap, name, None)
-        if isinstance(cls, type):
+        if isinstance(cls, type) and cls.__name__ != "frame":
             pickle(cls.__bases__[0], cls.__reduce__)
 
+    try:
+        # in case of reload(stackless)
+        reduce_frame = _wrap.reduce_frame
+    except AttributeError:
+        from weakref import WeakValueDictionary
+
+        wrap_set_reduce_frame = _wrap.set_reduce_frame
+
+        def set_reduce_frame(func):
+            wrap_set_reduce_frame(func)
+            _wrap.reduce_frame = func
+
+        _wrap.set_reduce_frame = set_reduce_frame
+
+        wrap_frame__reduce = _wrap.frame.__reduce__
+        cache = WeakValueDictionary()
+
+        class _Frame_Wrapper(object):
+            """Wrapper for frames to be pickled"""
+            __slots__ = ('__weakref__', 'frame')
+
+            @classmethod
+            def reduce_frame(cls, frame):
+                oid = id(frame)
+                try:
+                    return cache[oid]
+                except KeyError:
+                    cache[oid] = reducer = cls(frame)
+                    return reducer
+
+            def __init__(self, frame):
+                self.frame = frame
+
+            def __reduce__(self):
+                return wrap_frame__reduce(self.frame)
+        reduce_frame = _Frame_Wrapper.reduce_frame
+    _wrap.set_reduce_frame(reduce_frame)
+
+
     class StacklessModuleType(types.ModuleType):
 
         def current(self):

Stackless/core/stackless_impl.h

@@ -388,6 +388,8 @@ PyObject * slp_restore_exception(PyFrameObject *f, int exc, PyObject *retval);
 PyObject * slp_restore_tracing(PyFrameObject *f, int exc, PyObject *retval);
 /* other eval_frame functions from Objects/typeobject.c */
 PyObject * slp_tp_init_callback(PyFrameObject *f, int exc, PyObject *retval);
+/* functions related to pickling */
+PyObject * slp_reduce_frame(PyFrameObject * frame);
 
 /* rebirth of software stack avoidance */
 

Stackless/module/taskletobject.c

@@ -442,7 +442,14 @@ tasklet_reduce(PyTaskletObject * t)
                 goto err_exit;
         }
         if (append_frame) {
-            if (PyList_Append(lis, (PyObject *) f)) goto err_exit;
+            int ret;
+            PyObject * frame_reducer = slp_reduce_frame(f);
+            if (frame_reducer == NULL)
+                goto err_exit;
+            ret = PyList_Append(lis, frame_reducer);
+            Py_DECREF(frame_reducer);
+            if (ret)
+                goto err_exit;
         }
         f = f->f_back;
     }
@@ -1486,8 +1493,8 @@ tasklet_get_frame(PyTaskletObject *task)
 {
     PyObject *ret = (PyObject*) PyTasklet_GetFrame(task);
     if (ret)
-		return ret;
-	Py_RETURN_NONE;
+        return ret;
+    Py_RETURN_NONE;
 }
 
 PyObject *

Stackless/pickling/prickelpit.c

@@ -179,6 +179,36 @@ static struct _typeobject wrap_##type = { \
 };
 
 static PyObject *types_mod = NULL;
+static PyObject *reduce_frame_func = NULL;
+
+PyDoc_STRVAR(set_reduce_frame__doc__,
+"set_reduce_frame(func) -- set the function used to reduce frames during pickling.\n"
+"The function takes a frame as its sole argument and must return a pickleable object.\n");
+
+static PyObject *
+set_reduce_frame(PyObject *self, PyObject *func)
+{
+    if (func == Py_None) {
+        Py_CLEAR(reduce_frame_func);
+    } else {
+        if (!PyCallable_Check(func)) {
+            TYPE_ERROR("func must be callable", NULL);
+        }
+        Py_INCREF(func);
+        Py_XSETREF(reduce_frame_func, func);
+    }
+    Py_RETURN_NONE;
+}
+
+PyObject *
+slp_reduce_frame(PyFrameObject * frame) {
+    if (!PyFrame_Check(frame) || reduce_frame_func == NULL) {
+        Py_INCREF(frame);
+        return (PyObject *)frame;
+    }
+    return PyObject_CallFunctionObjArgs(reduce_frame_func, (PyObject *)frame, NULL);
+}
+
 
 static struct PyMethodDef _new_methoddef[] = {
     {"__new__", (PyCFunction)_new_wrapper, METH_VARARGS | METH_KEYWORDS,
@@ -1156,13 +1186,19 @@ static PyObject *
 tb_reduce(tracebackobject * tb)
 {
     PyObject *tup = NULL;
-    char *fmt = "(O()(OiiO))";
+    PyObject *frame_reducer;
+    const char *fmt = "(O()(OiiO))";
 
     if (tb->tb_next == NULL)
         fmt = "(O()(Oii))";
+    frame_reducer = slp_reduce_frame(tb->tb_frame);
+    if (frame_reducer == NULL)
+        return NULL;
+
     tup = Py_BuildValue(fmt,
                         &wrap_PyTraceBack_Type,
-                        tb->tb_frame, tb->tb_lasti, tb->tb_lineno, tb->tb_next);
+                        frame_reducer, tb->tb_lasti, tb->tb_lineno, tb->tb_next);
+    Py_DECREF(frame_reducer);
     return tup;
 }
 
@@ -2242,11 +2278,16 @@ static PyObject *
 gen_reduce(PyGenObject *gen)
 {
     PyObject *tup;
+    PyObject *frame_reducer;
+    frame_reducer = slp_reduce_frame(gen->gi_frame);
+    if (frame_reducer == NULL)
+        return NULL;
     tup = Py_BuildValue("(O()(Oi))",
                         &wrap_PyGen_Type,
-                        gen->gi_frame,
+                        frame_reducer,
                         gen->gi_running
                         );
+    Py_DECREF(frame_reducer);
     return tup;
 }
 
@@ -2458,22 +2499,29 @@ static int
 _wrapmodule_traverse(PyObject *self, visitproc visit, void *arg)
 {
     Py_VISIT(gen_exhausted_frame);
+    Py_VISIT(reduce_frame_func);
     return 0;
 }
 
 static int
 _wrapmodule_clear(PyObject *self)
 {
     Py_CLEAR(gen_exhausted_frame);
+    Py_CLEAR(reduce_frame_func);
     return 0;
 }
 
+static PyMethodDef _wrapmodule_methods[] = {
+    {"set_reduce_frame", set_reduce_frame, METH_O, set_reduce_frame__doc__},
+    {NULL,                          NULL}       /* sentinel */
+};
+
 static struct PyModuleDef _wrapmodule = {
     PyModuleDef_HEAD_INIT,
     "_stackless._wrap",
     NULL,
     -1,
-    NULL,
+    _wrapmodule_methods,
     NULL,
     _wrapmodule_traverse,
     _wrapmodule_clear,

Stackless/unittests/support.py

@@ -747,6 +747,14 @@ def helper():
         return result
 
 
+def get_reduce_frame():
+    """counterpart to stackless._wrap.set_reduce_frame()
+
+    Only for testing!
+    """
+    return getattr(stackless._wrap, "reduce_frame", None)
+
+
 def test_main():
     """Main function for the CPython :mod:`test.regrtest` test driver.
 

Stackless/unittests/test_defects.py

@@ -14,7 +14,8 @@
 
 from stackless import _test_nostacklesscall as apply_not_stackless
 from support import test_main  # @UnusedImport
-from support import StacklessTestCase, captured_stderr, require_one_thread
+from support import (StacklessTestCase, captured_stderr, require_one_thread,
+                     get_reduce_frame)
 
 
 """
@@ -236,6 +237,7 @@ def testCrasher(self):
         frameType = type(frame)
         while frame and frame.f_back:
             frame = frame.f_back
+        frame = get_reduce_frame()(frame)
         p = pickle.dumps(frame, -1)
         frame = None
         frame = pickle.loads(p)

Stackless/unittests/test_pickle.py

@@ -7,7 +7,7 @@
 from stackless import schedule, tasklet, stackless
 
 from support import test_main  # @UnusedImport
-from support import StacklessTestCase, StacklessPickleTestCase
+from support import StacklessTestCase, StacklessPickleTestCase, get_reduce_frame
 
 
 # because test runner instances in the testsuite contain copies of the old stdin/stdout thingies,
@@ -494,6 +494,14 @@ def testFunctionModulePreservation(self):
 
 
 class TestFramePickling(StacklessTestCase):
+    def test_get_set_reduce_frame(self):
+        # test setting / getting the reduce frame function
+        rf = get_reduce_frame()
+        self.assertTrue(callable(rf))
+        stackless._wrap.set_reduce_frame(None)
+        self.assertIsNone(get_reduce_frame())
+        stackless._wrap.set_reduce_frame(rf)
+        self.assertIs(get_reduce_frame(), rf)
 
     def testLocalplus(self):
         result = []
@@ -607,6 +615,7 @@ def d():
         p = self.dumps(tb)
         tb2 = self.loads(p)
         # basics
+        innerframes_orig = inspect.getinnerframes(tb)
         self.assertIs(type(tb), type(tb2))
         self.assertIsNot(tb, tb2)
         innerframes = inspect.getinnerframes(tb2)

Stackless/unittests/test_thread.py

@@ -235,6 +235,9 @@ def to_current_thread(self, task):
                     frameList[i] = newFrame
         # rebind the task
         task = reducedTask[0](*reducedTask[1])
+        for i in range(len(reducedTask[2][3])):
+            if not isinstance(reducedTask[2][3][i], stackless.cframe):
+                reducedTask[2][3][i] = reducedTask[2][3][i].frame
         task.__setstate__(reducedTask[2])
         return task
 

Stackless/unittests/test_tstate.py

@@ -6,7 +6,7 @@
 from stackless import *
 
 from support import test_main  # @UnusedImport
-from support import StacklessTestCase
+from support import StacklessTestCase, get_reduce_frame
 
 # import os
 # def debug():
@@ -332,9 +332,10 @@ def testReduceOfTracingState(self):
             return
 
         # test if the tracing cframe is present / not present
-        self.assertListEqual(reduced_tasklet1[2][3], [frame])
+        reduce_frame = get_reduce_frame()
+        self.assertListEqual(reduced_tasklet1[2][3], [reduce_frame(frame)])
         self.assertEquals(len(reduced_tasklet2[2][3]), 2)
-        self.assertIs(reduced_tasklet2[2][3][0], frame)
+        self.assertIs(reduced_tasklet2[2][3][0], reduce_frame(frame))
         self.assertIsInstance(reduced_tasklet2[2][3][1], stackless.cframe)
 
         cf = reduced_tasklet2[2][3][1]