Implement PEP 567 support (contextvars) for Python 3.7 (#155)

Author: 1st1

tests/test_context.py

@@ -0,0 +1,139 @@
+import asyncio
+import decimal
+import random
+import sys
+import unittest
+
+from uvloop import _testbase as tb
+
+
+PY37 = sys.version_info >= (3, 7, 0)
+
+
+class _ContextBaseTests:
+
+    @unittest.skipUnless(PY37, 'requires Python 3.7')
+    def test_task_decimal_context(self):
+        async def fractions(t, precision, x, y):
+            with decimal.localcontext() as ctx:
+                ctx.prec = precision
+                a = decimal.Decimal(x) / decimal.Decimal(y)
+                await asyncio.sleep(t, loop=self.loop)
+                b = decimal.Decimal(x) / decimal.Decimal(y ** 2)
+                return a, b
+
+        async def main():
+            r1, r2 = await asyncio.gather(
+                fractions(0.1, 3, 1, 3), fractions(0.2, 6, 1, 3),
+                loop=self.loop)
+
+            return r1, r2
+
+        r1, r2 = self.loop.run_until_complete(main())
+
+        self.assertEqual(str(r1[0]), '0.333')
+        self.assertEqual(str(r1[1]), '0.111')
+
+        self.assertEqual(str(r2[0]), '0.333333')
+        self.assertEqual(str(r2[1]), '0.111111')
+
+    @unittest.skipUnless(PY37, 'requires Python 3.7')
+    def test_task_context_1(self):
+        import contextvars
+        cvar = contextvars.ContextVar('cvar', default='nope')
+
+        async def sub():
+            await asyncio.sleep(0.01, loop=self.loop)
+            self.assertEqual(cvar.get(), 'nope')
+            cvar.set('something else')
+
+        async def main():
+            self.assertEqual(cvar.get(), 'nope')
+            subtask = self.loop.create_task(sub())
+            cvar.set('yes')
+            self.assertEqual(cvar.get(), 'yes')
+            await subtask
+            self.assertEqual(cvar.get(), 'yes')
+
+        task = self.loop.create_task(main())
+        self.loop.run_until_complete(task)
+
+    @unittest.skipUnless(PY37, 'requires Python 3.7')
+    def test_task_context_2(self):
+        import contextvars
+        cvar = contextvars.ContextVar('cvar', default='nope')
+
+        async def main():
+            def fut_on_done(fut):
+                # This change must not pollute the context
+                # of the "main()" task.
+                cvar.set('something else')
+
+            self.assertEqual(cvar.get(), 'nope')
+
+            for j in range(2):
+                fut = self.loop.create_future()
+                fut.add_done_callback(fut_on_done)
+                cvar.set('yes{}'.format(j))
+                self.loop.call_soon(fut.set_result, None)
+                await fut
+                self.assertEqual(cvar.get(), 'yes{}'.format(j))
+
+                for i in range(3):
+                    # Test that task passed its context to add_done_callback:
+                    cvar.set('yes{}-{}'.format(i, j))
+                    await asyncio.sleep(0.001, loop=self.loop)
+                    self.assertEqual(cvar.get(), 'yes{}-{}'.format(i, j))
+
+        task = self.loop.create_task(main())
+        self.loop.run_until_complete(task)
+
+        self.assertEqual(cvar.get(), 'nope')
+
+    @unittest.skipUnless(PY37, 'requires Python 3.7')
+    def test_task_context_3(self):
+        import contextvars
+        cvar = contextvars.ContextVar('cvar', default=-1)
+
+        # Run 100 Tasks in parallel, each modifying cvar.
+
+        async def sub(num):
+            for i in range(10):
+                cvar.set(num + i)
+                await asyncio.sleep(
+                    random.uniform(0.001, 0.05), loop=self.loop)
+                self.assertEqual(cvar.get(), num + i)
+
+        async def main():
+            tasks = []
+            for i in range(100):
+                task = self.loop.create_task(sub(random.randint(0, 10)))
+                tasks.append(task)
+
+            await asyncio.gather(
+                *tasks, loop=self.loop, return_exceptions=True)
+
+        self.loop.run_until_complete(main())
+
+        self.assertEqual(cvar.get(), -1)
+
+
+class Test_UV_Context(_ContextBaseTests, tb.UVTestCase):
+
+    @unittest.skipIf(PY37, 'requires Python <3.6')
+    def test_context_arg(self):
+        def cb():
+            pass
+
+        with self.assertRaisesRegex(RuntimeError, 'requires Python 3.7'):
+            self.loop.call_soon(cb, context=1)
+
+        with self.assertRaisesRegex(RuntimeError, 'requires Python 3.7'):
+            self.loop.call_soon_threadsafe(cb, context=1)
+
+        with self.assertRaisesRegex(RuntimeError, 'requires Python 3.7'):
+            self.loop.call_later(0.1, cb, context=1)
+
+
+class Test_AIO_Context(_ContextBaseTests, tb.AIOTestCase):
+    pass

uvloop/cbhandles.pxd

@@ -1,6 +1,7 @@
 cdef class Handle:
     cdef:
         Loop loop
+        object context
         bint _cancelled
 
         str meth_name
@@ -13,6 +14,8 @@ cdef class Handle:
         readonly _source_traceback
 
     cdef inline _set_loop(self, Loop loop)
+    cdef inline _set_context(self, object context)
+
     cdef inline _run(self)
     cdef _cancel(self)
 
@@ -25,6 +28,7 @@ cdef class TimerHandle:
         bint _cancelled
         UVTimer timer
         Loop loop
+        object context
         object __weakref__
 
         readonly _source_traceback

uvloop/cbhandles.pyx

@@ -14,6 +14,16 @@ cdef class Handle:
         if loop._debug:
             self._source_traceback = extract_stack()
 
+    cdef inline _set_context(self, object context):
+        if PY37:
+            if context is None:
+                context = <object>PyContext_CopyCurrent()
+            self.context = context
+        else:
+            if context is not None:
+                raise RuntimeError('"context" argument requires Python 3.7')
+            self.context = None
+
     def __dealloc__(self):
         if UVLOOP_DEBUG and self.loop is not None:
             self.loop._debug_cb_handles_count -= 1
@@ -39,6 +49,10 @@ cdef class Handle:
                           # we guard 'self' manually (since the callback
                           # might cause GC of the handle.)
         try:
+            if PY37:
+                assert self.context is not None
+                PyContext_Enter(<PyContext*>self.context)
+
             if cb_type == 1:
                 callback = self.arg1
                 args = self.arg2
@@ -83,7 +97,10 @@ cdef class Handle:
             self.loop.call_exception_handler(context)
 
         finally:
+            context = self.context
             Py_DECREF(self)
+            if PY37:
+                PyContext_Exit(<PyContext*>context)
 
     cdef _cancel(self):
         self._cancelled = 1
@@ -148,13 +165,26 @@ cdef class Handle:
 @cython.freelist(DEFAULT_FREELIST_SIZE)
 cdef class TimerHandle:
     def __cinit__(self, Loop loop, object callback, object args,
-                  uint64_t delay):
+                  uint64_t delay, object context):
 
         self.loop = loop
         self.callback = callback
         self.args = args
         self._cancelled = 0
 
+        if UVLOOP_DEBUG:
+            self.loop._debug_cb_timer_handles_total += 1
+            self.loop._debug_cb_timer_handles_count += 1
+
+        if PY37:
+            if context is None:
+                context = <object>PyContext_CopyCurrent()
+            self.context = context
+        else:
+            if context is not None:
+                raise RuntimeError('"context" argument requires Python 3.7')
+            self.context = None
+
         if loop._debug:
             self._source_traceback = extract_stack()
 
@@ -166,10 +196,6 @@ cdef class TimerHandle:
         # Only add to loop._timers when `self.timer` is successfully created
         loop._timers.add(self)
 
-        if UVLOOP_DEBUG:
-            self.loop._debug_cb_timer_handles_total += 1
-            self.loop._debug_cb_timer_handles_count += 1
-
     def __dealloc__(self):
         if UVLOOP_DEBUG:
             self.loop._debug_cb_timer_handles_count -= 1
@@ -207,6 +233,10 @@ cdef class TimerHandle:
         if self.loop._debug:
             started = time_monotonic()
         try:
+            if PY37:
+                assert self.context is not None
+                PyContext_Enter(<PyContext*>self.context)
+
             if args is not None:
                 callback(*args)
             else:
@@ -230,7 +260,10 @@ cdef class TimerHandle:
                         'Executing %r took %.3f seconds',
                         self, delta)
         finally:
+            context = self.context
             Py_DECREF(self)
+            if PY37:
+                PyContext_Exit(<PyContext*>context)
 
     # Public API
 
@@ -263,10 +296,12 @@ cdef class TimerHandle:
         self._cancel()
 
 
-cdef new_Handle(Loop loop, object callback, object args):
+
+cdef new_Handle(Loop loop, object callback, object args, object context):
     cdef Handle handle
     handle = Handle.__new__(Handle)
     handle._set_loop(loop)
+    handle._set_context(context)
 
     handle.cb_type = 1
 
@@ -280,6 +315,7 @@ cdef new_MethodHandle(Loop loop, str name, method_t callback, object ctx):
     cdef Handle handle
     handle = Handle.__new__(Handle)
     handle._set_loop(loop)
+    handle._set_context(None)
 
     handle.cb_type = 2
     handle.meth_name = name
@@ -296,6 +332,7 @@ cdef new_MethodHandle1(Loop loop, str name, method1_t callback,
     cdef Handle handle
     handle = Handle.__new__(Handle)
     handle._set_loop(loop)
+    handle._set_context(None)
 
     handle.cb_type = 3
     handle.meth_name = name
@@ -312,6 +349,7 @@ cdef new_MethodHandle2(Loop loop, str name, method2_t callback, object ctx,
     cdef Handle handle
     handle = Handle.__new__(Handle)
     handle._set_loop(loop)
+    handle._set_context(None)
 
     handle.cb_type = 4
     handle.meth_name = name
@@ -329,6 +367,7 @@ cdef new_MethodHandle3(Loop loop, str name, method3_t callback, object ctx,
     cdef Handle handle
     handle = Handle.__new__(Handle)
     handle._set_loop(loop)
+    handle._set_context(None)
 
     handle.cb_type = 5
     handle.meth_name = name

uvloop/includes/compat.h

@@ -1,4 +1,5 @@
 #include <errno.h>
+#include "Python.h"
 #include "uv.h"
 
 
@@ -27,3 +28,22 @@
 struct epoll_event {};
 int epoll_ctl(int epfd, int op, int fd, struct epoll_event *event) {};
 #endif
+
+
+#if PY_VERSION_HEX < 0x03070000
+typedef struct {
+    PyObject_HEAD
+} PyContext;
+
+PyContext * PyContext_CopyCurrent(void) {
+    return NULL;
+};
+
+int PyContext_Enter(PyContext *ctx) {
+    return -1;
+}
+
+int PyContext_Exit(PyContext *ctx) {
+    return -1;
+}
+#endif

uvloop/includes/python.pxd

@@ -1,7 +1,9 @@
 cdef extern from "Python.h":
+    int PY_VERSION_HEX
+
     void* PyMem_RawMalloc(size_t n)
     void* PyMem_RawRealloc(void *p, size_t n)
-    void* PyMem_RawCalloc(size_t nelem, size_t elsize)  # Python >= 3.5!
+    void* PyMem_RawCalloc(size_t nelem, size_t elsize)
     void PyMem_RawFree(void *p)
 
     object PyUnicode_EncodeFSDefault(object)
@@ -11,3 +13,10 @@ cdef extern from "Python.h":
     void _PyImport_AcquireLock()
     int _PyImport_ReleaseLock()
     void _Py_RestoreSignals()
+
+
+cdef extern from "includes/compat.h":
+    ctypedef struct PyContext
+    PyContext* PyContext_CopyCurrent() except NULL
+    int PyContext_Enter(PyContext *) except -1
+    int PyContext_Exit(PyContext *) except -1

uvloop/loop.pxd

@@ -142,10 +142,11 @@ cdef class Loop:
     cdef inline _queue_write(self, UVStream stream)
     cdef _exec_queued_writes(self)
 
-    cdef inline _call_soon(self, object callback, object args)
+    cdef inline _call_soon(self, object callback, object args, object context)
     cdef inline _call_soon_handle(self, Handle handle)
 
-    cdef _call_later(self, uint64_t delay, object callback, object args)
+    cdef _call_later(self, uint64_t delay, object callback, object args,
+                     object context)
 
     cdef void _handle_exception(self, object ex)