Patch call_soon and friends.

README.rst

@@ -14,6 +14,16 @@ experimentally provides the missing asyncio support for the
 ``contextvars`` backport library. Please read more in Python 3.7 `contextvars
 documentation <https://docs.python.org/3/library/contextvars.html>`_.
 
+Usage
+-----
+
+Import the `aiocontextvars` package before you create any aio tasks or loops::
+
+    import aiocontextvars  # noqa: F401
+
+You can then import from the usual `contextvars` package to use the contextvars::
+
+    from contextvars import ContextVar
 
 Compatibility
 -------------
@@ -28,10 +38,13 @@ Python 3.7 ``contextvars`` implementation:
 
 Python 3.7 added keyword argument ``context`` to ``call_soon()`` and its family
 methods. By default those methods will copy (inherit) the current context and
-run the given method in that context. But ``aiocontextvars`` won't touch the
-loop, so in order to achieve the same effect, you'll need to::
+run the given method in that context. 
+
+By default ``aiocontextvars`` will copy the current context with the patched
+function, however it won't recognise the context parameter. If you need to 
+run under a particular context, not copied from the current context use::
 
-    loop.call_soon(copy_context().run, my_meth)
+    loop.call_soon(context.run, my_meth)
 
 2. Task local.
 

aiocontextvars.py

@@ -23,51 +23,71 @@ def _get_context():
             state.context = ctx
         return ctx
 
-
     def _set_context(ctx):
         state = _get_state()
         state.context = ctx
 
-
     def _get_state():
         loop = asyncio._get_running_loop()
         if loop is None:
             return _state
         task = asyncio.Task.current_task(loop=loop)
         return _state if task is None else task
 
-
     contextvars._get_context = _get_context
     contextvars._set_context = _set_context
 
-
     def create_task(loop, coro):
         task = loop._orig_create_task(coro)
         if task._source_traceback:
             del task._source_traceback[-1]
         task.context = copy_context()
         return task
 
+    def call_soon(loop, callback, *args):
+        if callback == Context.run:
+            return loop._orig_call_soon(callback, *args)
+        else:
+            context = copy_context()
+            return loop._orig_call_soon(context.run, lambda: callback(*args))
+
+    def call_later(loop, delay, callback, *args):
+        if callback == Context.run:
+            return loop._orig_call_later(delay, callback, *args)
+        else:
+            context = copy_context()
+            return loop._orig_call_later(
+                delay, context.run, lambda: callback(*args))
+
+    def call_at(loop, when, callback, *args):
+        if callback == Context.run:
+            return loop._orig_call_at(when, callback, *args)
+        else:
+            context = copy_context()
+            return loop._orig_call_at(
+                when, context.run, lambda: callback(*args))
 
     def _patch_loop(loop):
         if loop and not hasattr(loop, '_orig_create_task'):
             loop._orig_create_task = loop.create_task
             loop.create_task = types.MethodType(create_task, loop)
+            loop._orig_call_soon = loop.call_soon
+            loop.call_soon = types.MethodType(call_soon, loop)
+            loop._orig_call_later = loop.call_later
+            loop.call_later = types.MethodType(call_later, loop)
+            loop._orig_call_at = loop.call_at
+            loop.call_at = types.MethodType(call_at, loop)
         return loop
 
-
     def get_event_loop():
         return _patch_loop(_get_event_loop())
 
-
     def set_event_loop(loop):
         return _set_event_loop(_patch_loop(loop))
 
-
     def new_event_loop():
         return _patch_loop(_new_event_loop())
 
-
     _get_event_loop = asyncio.get_event_loop
     _set_event_loop = asyncio.set_event_loop
     _new_event_loop = asyncio.new_event_loop

tests/test_aio.py

@@ -1,6 +1,9 @@
 import asyncio
 import pytest
-from aiocontextvars import ContextVar
+from asyncio import Event
+from contextvars import ContextVar
+import aiocontextvars  # noqa: F401
+
 
 v = ContextVar('v')
 
@@ -38,6 +41,52 @@ async def sub():
     assert v.get('default') == 'update'
 
 
+# noinspection PyUnusedLocal
+@pytest.mark.asyncio
+async def test_inherit_call_soon():
+    v.set('initial')
+
+    event = Event()
+
+    def sub(arg1, arg2):
+        try:
+            assert v.get('default') == 'initial'
+            assert arg1 == 'arg1'
+            assert arg2 == 'arg2'
+            v.set('sub')
+        finally:
+            event.set()
+
+    loop = asyncio.get_event_loop()
+    loop.call_soon(sub, 'arg1', 'arg2')
+    v.set('update')
+
+    await asyncio.sleep(0.1)
+    await event.wait()
+    assert v.get('default') == 'update'
+
+# noinspection PyUnusedLocal
+@pytest.mark.asyncio
+async def test_inherit_call_later():
+    v.set('initial')
+
+    event = Event()
+
+    def sub():
+        try:
+            assert v.get('default') == 'initial'
+            v.set('sub')
+        finally:
+            event.set()
+
+    loop = asyncio.get_event_loop()
+    loop.call_later(0.5, sub)
+    v.set('update')
+
+    await event.wait()
+    assert v.get('default') == 'update'
+
+
 def test_set_event_loop_none():
     # Setting event loop to None should not raise an exception
     asyncio.set_event_loop(None)

tests/test_var.py

@@ -1,5 +1,6 @@
 import pytest
-from aiocontextvars import ContextVar
+from contextvars import ContextVar
+import aiocontextvars  # noqa: F401
 
 
 def test_get():