Stackless issue #197: stackless call method "contextvars.Context.run"

Enable stackless calls of method "contextvars.Context.run", if soft-switching is enabled.

Author: Anselm Kruis

Python/context.c

@@ -4,6 +4,8 @@
 #include "internal/pystate.h"
 #include "internal/context.h"
 #include "internal/hamt.h"
+#include "core/stackless_impl.h"
+#include "pickling/prickelpit.h"
 
 
 #define CONTEXT_FREELIST_MAXLEN 255
@@ -562,10 +564,33 @@ _contextvars_Context_copy_impl(PyContext *self)
 }
 
 
+#ifdef STACKLESS
+static PyObject* context_run_callback(PyFrameObject *f, int exc, PyObject *result)
+{
+    PyCFrameObject *cf = (PyCFrameObject *)f;
+    assert(PyContext_CheckExact(cf->ob1));
+    PyContext *context = (PyContext *)cf->ob1;
+    cf->ob1 = NULL;
+
+    if (PyContext_Exit(context)) {
+        Py_CLEAR(result);
+    }
+
+    Py_DECREF(context);
+    SLP_STORE_NEXT_FRAME(PyThreadState_GET(), cf->f_back);
+    return result;
+}
+
+SLP_DEF_INVALID_EXEC(context_run_callback)
+#endif
+
+
 static PyObject *
 context_run(PyContext *self, PyObject *const *args,
             Py_ssize_t nargs, PyObject *kwnames)
 {
+    STACKLESS_GETARG();
+
     if (nargs < 1) {
         PyErr_SetString(PyExc_TypeError,
                         "run() missing 1 required positional argument");
@@ -576,9 +601,42 @@ context_run(PyContext *self, PyObject *const *args,
         return NULL;
     }
 
+#ifdef STACKLESS
+    PyThreadState *ts = PyThreadState_GET();
+    PyCFrameObject *f = NULL;
+    if (stackless) {
+        f = slp_cframe_new(context_run_callback, 1);
+        if (f == NULL)
+            return NULL;
+        Py_INCREF(self);
+        f->ob1 = (PyObject *)self;
+        SLP_SET_CURRENT_FRAME(ts, (PyFrameObject *)f);
+        /* f contains the only counted reference to current frame. This reference
+         * keeps the fame alive during the following _PyObject_FastCallKeywords().
+         */
+    }
+#endif
+    STACKLESS_PROMOTE_ALL();
+
     PyObject *call_result = _PyObject_FastCallKeywords(
         args[0], args + 1, nargs - 1, kwnames);
 
+    STACKLESS_ASSERT();
+#ifdef STACKLESS
+    if (stackless && !STACKLESS_UNWINDING(call_result)) {
+        /* required, because we added a C-frame */
+        assert(f);
+        assert((PyFrameObject *)f == SLP_CURRENT_FRAME(ts));
+        SLP_STORE_NEXT_FRAME(ts, (PyFrameObject *)f);
+        Py_DECREF(f);
+        return STACKLESS_PACK(ts, call_result);
+    }
+    Py_XDECREF(f);
+    if (STACKLESS_UNWINDING(call_result)) {
+        return call_result;
+    }
+#endif
+
     if (PyContext_Exit(self)) {
         return NULL;
     }
@@ -593,7 +651,7 @@ static PyMethodDef PyContext_methods[] = {
     _CONTEXTVARS_CONTEXT_KEYS_METHODDEF
     _CONTEXTVARS_CONTEXT_VALUES_METHODDEF
     _CONTEXTVARS_CONTEXT_COPY_METHODDEF
-    {"run", (PyCFunction)context_run, METH_FASTCALL | METH_KEYWORDS, NULL},
+    {"run", (PyCFunction)context_run, METH_FASTCALL | METH_KEYWORDS | METH_STACKLESS, NULL},
     {NULL, NULL}
 };
 
@@ -1222,5 +1280,13 @@ _PyContext_Init(void)
     }
     Py_DECREF(missing);
 
+#ifdef STACKLESS
+    if (slp_register_execute(&PyCFrame_Type, "context_run_callback",
+        context_run_callback, SLP_REF_INVALID_EXEC(context_run_callback)) != 0)
+    {
+        return 0;
+    }
+#endif
+
     return 1;
 }

Stackless/changelog.txt

@@ -9,6 +9,10 @@ What's New in Stackless 3.X.X?
 
 *Release date: 20XX-XX-XX*
 
+- https://github.com/stackless-dev/stackless/issues/197
+  Enable stackless calls of method "contextvars.Context.run", if soft-switching
+  is enabled.
+
 - https://github.com/stackless-dev/stackless/issues/188
   Enable stackless calls of sub-iterators and coroutines (technically: the
   YIELD_FROM opcode) and of the following methods, if soft-switching is

Stackless/unittests/test_generator.py

@@ -6,6 +6,8 @@
 import pickle
 import contextlib
 import sys
+import contextvars
+import asyncio
 
 from support import test_main  # @UnusedImport
 from support import StacklessTestCase, captured_stderr
@@ -361,6 +363,27 @@ async def test():
         n = self.run_until_complete(t)
         self.assertEqual(n, 0)
 
+    def test_context_run(self):
+        contextvars.Context().run(self.assertLevel)
+
+    #  needs Stackless pull request #188
+    def xx_test_asyncio(self):
+        async def test():
+            try:
+                await self.coro()
+            finally:
+                loop.stop()
+
+        asyncio.set_event_loop(asyncio.new_event_loop())
+        self.addCleanup(asyncio.set_event_loop, None)
+        loop = asyncio.get_event_loop()
+        task = asyncio.tasks._PyTask(test())
+        asyncio.ensure_future(task)
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
 
 if __name__ == '__main__':
     unittest.main()