Support context in loop.create_task()

Author: fantix

tests/test_base.py

@@ -540,7 +540,9 @@ class MyTask(asyncio.Task):
         async def coro():
             pass
 
-        factory = lambda loop, coro: MyTask(coro, loop=loop)
+        factory = lambda loop, coro, **kwargs: MyTask(
+            coro, loop=loop, **kwargs
+        )
 
         self.assertIsNone(self.loop.get_task_factory())
         self.loop.set_task_factory(factory)
@@ -577,7 +579,9 @@ def get_name(self):
         async def coro():
             pass
 
-        factory = lambda loop, coro: MyTask(coro, loop=loop)
+        factory = lambda loop, coro, **kwargs: MyTask(
+            coro, loop=loop, **kwargs
+        )
 
         self.assertIsNone(self.loop.get_task_factory())
         task = self.loop.create_task(coro(), name="mytask")

uvloop/loop.pyx

@@ -50,6 +50,7 @@ include "errors.pyx"
 
 cdef:
     int PY39 = PY_VERSION_HEX >= 0x03090000
+    int PY311 = PY_VERSION_HEX >= 0x030b0000
     uint64_t MAX_SLEEP = 3600 * 24 * 365 * 100
 
 
@@ -1413,19 +1414,31 @@ cdef class Loop:
         """Create a Future object attached to the loop."""
         return self._new_future()
 
-    def create_task(self, coro, *, name=None):
+    def create_task(self, coro, *, name=None, context=None):
         """Schedule a coroutine object.
 
         Return a task object.
 
         If name is not None, task.set_name(name) will be called if the task
         object has the set_name attribute, true for default Task in Python 3.8.
+
+        An optional keyword-only context argument allows specifying a custom
+        contextvars.Context for the coro to run in. The current context copy is
+        created when no context is provided. Only available in Python 3.11.
         """
         self._check_closed()
-        if self._task_factory is None:
-            task = aio_Task(coro, loop=self)
+        if PY311:
+            if self._task_factory is None:
+                task = aio_Task(coro, loop=self, context=context)
+            else:
+                task = self._task_factory(self, coro, context=context)
         else:
-            task = self._task_factory(self, coro)
+            if UVLOOP_DEBUG:
+                assert context is None
+            if self._task_factory is None:
+                task = aio_Task(coro, loop=self)
+            else:
+                task = self._task_factory(self, coro)
 
         # copied from asyncio.tasks._set_task_name (bpo-34270)
         if name is not None: