diff --git a/tests/test_context.py b/tests/test_context.py index c57e359..c167abe 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -3,6 +3,7 @@ import decimal import random import sys import unittest +import weakref from uvloop import _testbase as tb @@ -117,6 +118,29 @@ class _ContextBaseTests: self.assertEqual(cvar.get(), -1) + @unittest.skipUnless(PY37, 'requires Python 3.7') + def test_task_context_4(self): + import contextvars + cvar = contextvars.ContextVar('cvar', default='nope') + + class TrackMe: + pass + tracked = TrackMe() + ref = weakref.ref(tracked) + + async def sub(): + cvar.set(tracked) + self.loop.call_soon(lambda: None) + + async def main(): + await self.loop.create_task(sub()) + + task = self.loop.create_task(main()) + self.loop.run_until_complete(task) + + del tracked + self.assertIsNone(ref()) + class Test_UV_Context(_ContextBaseTests, tb.UVTestCase): diff --git a/uvloop/cbhandles.pyx b/uvloop/cbhandles.pyx index c342064..031a7ec 100644 --- a/uvloop/cbhandles.pyx +++ b/uvloop/cbhandles.pyx @@ -15,9 +15,11 @@ cdef class Handle: self._source_traceback = extract_stack() cdef inline _set_context(self, object context): + cdef PyContext* current_context + if PY37: if context is None: - context = PyContext_CopyCurrent() + context = copy_current_context() self.context = context else: if context is not None: @@ -179,7 +181,7 @@ cdef class TimerHandle: if PY37: if context is None: - context = PyContext_CopyCurrent() + context = copy_current_context() self.context = context else: if context is not None: @@ -400,3 +402,15 @@ cdef extract_stack(): stack.reverse() return stack + + +cdef copy_current_context(): + cdef PyContext* current_context + + if PY37: + current_context = PyContext_CopyCurrent() + py_context = current_context + Py_XDECREF(current_context) + return py_context + + raise NotImplementedError('"contextvars" support requires Python 3.7+')